Keras Core is a preview release of Keras 3.0. Keras 3.0 will be generally available in Fall 2023
and will replace
Keras Core makes it possible to run Keras workflows on top of TensorFlow, JAX, and PyTorch. It also enables you to seamlessly integrate Keras components (like layers, models, or metrics) as part of low-level TensorFlow, JAX, and PyTorch workflows.
Module, can be exported as a TensorFlow
SavedModel, or can be instantiated as a stateless JAX function. That means that you can use your Keras Core models with PyTorch ecosystem packages, with the full range of TensorFlow deployment & production tools, and with JAX large-scale TPU training infrastructure. Write one
model.pyusing Keras Core APIs, and get access to everything the ML world has to offer.
predict()routines are compatible with
tf.data.Datasetobjects, with PyTorch
DataLoaderobjects, with NumPy arrays, Pandas dataframes -- regardless of the backend you're using. You can train a Keras Core + TensorFlow model on a PyTorch
DataLoaderor train a Keras Core + PyTorch model on a
You can install Keras Core via
pip install keras-core
You can then import it in Python:
import keras_core as keras
Note that Keras Core requires
tensorflow to be installed, as it uses the
tf.nest Python datastructure
preprocessing utility. In the future, this dependency will be removed, and you will only need
to install the specific backend frameworks you intend to use.
Note on cuDNN: If you intend to use Keras Core on GPU with multiple frameworks in the same environment, be mindful that the latest versions of the three backend frameworks tend to each require a different cuDNN version. Hence you will need to find a combination of framework versions that all work with the same cuDNN version (meanwhile, Keras Core is expected to work with both the latest and prior version of each backend framework, which gives you flexibility).
You can find the right combination of backend framework versions and cuDNN version by looking at what's installed by default on Colab, since Colab faces the same version compatibility issue.
To configure which backend the
keras-core package should use, can export the environment variable
or you can edit your local config file at
~/.keras/keras.json (it gets automatically created when you import
Available backend options are:
$ export KERAS_BACKEND="jax" $ python train.py
$ KERAS_BACKEND=jax python train.py
In Colab, you can use:
import os os.environ["KERAS_BACKEND"] = "jax"
Note that the backend must be configured before importing
and the backend cannot be changed after the package has been imported.
If you always want to use the same backend, just edit
specify your default backend.
As of version
0.6.0, KerasCV and KerasNLP support multiple backends with Keras
Core out of the box. There are two ways to configure these libraries to run with
multi-backend support. Using KerasCV as an example:
KERAS_BACKENDenvironment variable. If set, then KerasCV will be using Keras Core with the backend specified (e.g.,
.keras/keras_cv.jsonconfig files (which are automatically created the first time you import KerasCV):
Once that configuration step is done, you can just import KerasCV and start using it on top of your backend of choice:
import keras_cv import keras_core as keras filepath = keras.utils.get_file(origin="https://i.imgur.com/gCNcJJI.jpg") image = np.array(keras.utils.load_img(filepath)) image_resized = ops.image.resize(image, (640, 640))[None, ...] model = keras_cv.models.YOLOV8Detector.from_preset( "yolo_v8_m_pascalvoc", bounding_box_format="xywh", ) predictions = model.predict(image_resized)
KerasNLP works the same way once configured with
import keras_nlp gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset("gpt2_base_en") gpt2_lm.generate("My trip to Yosemite was", max_length=200)
Until Keras Core is officially released as Keras 3.0, KerasCV and KerasNLP will
tf.keras as the default backend. To restore this default behavior, simply
unset KERAS_BACKEND and ensure that
"multi_backend": False or is unset in
.keras/keras_nlp.json. You will need to restart the
Python runtime for changes to take effect.