Skip to main content

Multi-backend Keras.

Project description

Keras 3: A new multi-backend Keras

Keras 3 is a new multi-backend implementation of the Keras API, with support for TensorFlow, JAX, and PyTorch.

Installation

Install with pip

Keras 3 is available on PyPI as keras. Note that Keras 2 remains available as the tf-keras package.

  1. Install keras:
pip install keras --upgrade
  1. Install backend package(s).

To use keras, you should also install the backend of choice: tensorflow, jax, or torch. Note that tensorflow is required for using certain Keras 3 features: certain preprocessing layers as well as tf.data pipelines.

Local installation

Minimal installation

Keras 3 is compatible with Linux and MacOS systems. For Windows users, we recommend using WSL2 to run Keras. To install a local development version:

  1. Install dependencies:
pip install -r requirements.txt
  1. Run installation command from the root directory.
python pip_build.py --install

Adding GPU support

The requirements.txt file will install a CPU-only version of TensorFlow, JAX, and PyTorch. For GPU support, we also provide a separate requirements-{backend}-cuda.txt for TensorFlow, JAX, and PyTorch. These install all CUDA dependencies via pip and expect a NVIDIA driver to be pre-installed. We recommend a clean python environment for each backend to avoid CUDA version mismatches. As an example, here is how to create a Jax GPU environment with conda:

conda create -y -n keras-jax python=3.10
conda activate keras-jax
pip install -r requirements-jax-cuda.txt
python pip_build.py --install

Configuring your backend

You can export the environment variable KERAS_BACKEND or you can edit your local config file at ~/.keras/keras.json to configure your backend. Available backend options are: "tensorflow", "jax", "torch". Example:

export KERAS_BACKEND="jax"

In Colab, you can do:

import os
os.environ["KERAS_BACKEND"] = "jax"

import keras

Note: The backend must be configured before importing keras, and the backend cannot be changed after the package has been imported.

Backwards compatibility

Keras 3 is intended to work as a drop-in replacement for tf.keras (when using the TensorFlow backend). Just take your existing tf.keras code, make sure that your calls to model.save() are using the up-to-date .keras format, and you're done.

If your tf.keras model does not include custom components, you can start running it on top of JAX or PyTorch immediately.

If it does include custom components (e.g. custom layers or a custom train_step()), it is usually possible to convert it to a backend-agnostic implementation in just a few minutes.

In addition, Keras models can consume datasets in any format, regardless of the backend you're using: you can train your models with your existing tf.data.Dataset pipelines or PyTorch DataLoaders.

Why use Keras 3?

  • Run your high-level Keras workflows on top of any framework -- benefiting at will from the advantages of each framework, e.g. the scalability and performance of JAX or the production ecosystem options of TensorFlow.
  • Write custom components (e.g. layers, models, metrics) that you can use in low-level workflows in any framework.
    • You can take a Keras model and train it in a training loop written from scratch in native TF, JAX, or PyTorch.
    • You can take a Keras model and use it as part of a PyTorch-native Module or as part of a JAX-native model function.
  • Make your ML code future-proof by avoiding framework lock-in.
  • As a PyTorch user: get access to power and usability of Keras, at last!
  • As a JAX user: get access to a fully-featured, battle-tested, well-documented modeling and training library.

Read more in the Keras 3 release announcement.

Project details


Release history Release notifications | RSS feed

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

keras-nightly-3.0.3.dev2024011703.tar.gz (737.1 kB view details)

Uploaded Source

Built Distribution

File details

Details for the file keras-nightly-3.0.3.dev2024011703.tar.gz.

File metadata

File hashes

Hashes for keras-nightly-3.0.3.dev2024011703.tar.gz
Algorithm Hash digest
SHA256 c1796c8c9353f9fa1f3b10b91440a21d29501504749c13b93b94ae6e7d796026
MD5 9ef5733af480971d66b7cb294fa8cbe7
BLAKE2b-256 5acd27e8d1c1e1c403bb62f43b26b994c255302a0324750580b7d234652655fe

See more details on using hashes here.

File details

Details for the file keras_nightly-3.0.3.dev2024011703-py3-none-any.whl.

File metadata

File hashes

Hashes for keras_nightly-3.0.3.dev2024011703-py3-none-any.whl
Algorithm Hash digest
SHA256 0397facbe097637bea78c053281b91766fc6c1a68dbd70045da1c025723243d7
MD5 714187155df5d7c56a863ac675a77795
BLAKE2b-256 240691beaff73ee4a2b7f5d1d436b390497c46ac1fcbbea5d78045f417641913

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page