Multi-backend Keras.
Keras 3 is a new multi-backend implementation of the Keras API, with support for TensorFlow, JAX, and PyTorch.
Keras 3 is available on PyPI as keras
. Note that Keras 2 remains available as the tf-keras
package.
keras
:pip install keras --upgrade
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.
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:
pip install -r requirements.txt
python pip_build.py --install
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
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.
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
.
Module
or as part of a JAX-native model function.Read more in the Keras 3 release announcement.