keras

Keras 3: Deep Learning for Humans

Keras 3 is a sophisticated and versatile deep learning framework that supports JAX, TensorFlow, and PyTorch. It provides an effortless way to build and train models in various domains such as computer vision, natural language processing, audio processing, timeseries forecasting, and recommender systems. This flexibility makes it an ideal choice for both startups and large enterprises alike, as it efficiently scales from individual laptops to extensive clusters of GPUs or TPUs.

Key Features

• Accelerated Model Development: Developers can quickly ship deep learning solutions thanks to the high-level user experience provided by Keras. The framework also offers easy-to-debug runtimes like PyTorch and JAX eager execution, making the development process smoother and more efficient.
• State-of-the-Art Performance: By selecting the most efficient backend for a particular model architecture—often JAX—users can achieve significant speed enhancements, ranging from 20% to 350% compared to other frameworks.
• Datacenter-Scale Training: Keras 3 offers a seamless transition from local development environments to large-scale cloud environments, allowing developers to leverage advanced computing resources without extensive reconfiguration.

Installation Guide

Pip Installation

Keras 3 is conveniently available on PyPI under the package name keras. To get started with Keras 3:

1. Install Keras:

   pip install keras --upgrade
   

2. Install the required backend package(s):

   You need to choose and install one of the backend frameworks (tensorflow, jax, or torch). Note that some features of Keras 3, including specific preprocessing layers and tf.data pipelines, require TensorFlow.

Local Installation

For those developing locally, Keras 3 supports Linux and macOS systems, with Windows users advised to utilize WSL2 for optimal functionality.

Minimal Installation

1. Install necessary dependencies from the requirements.txt file:

   pip install -r requirements.txt
   

2. Execute the installation command from the project's root directory:

   python pip_build.py --install
   

3. If contributing to the codebase and updating public APIs, run the API generation script:

   ./shell/api_gen.sh
   

GPU Support

To take advantage of GPU acceleration, additional CUDA dependencies must be installed using requirements-{backend}-cuda.txt, tailored for each backend framework. Here is how to set up 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

Backend Configuration

The backend of Keras can be configured by setting the KERAS_BACKEND environment variable or by editing the ~/.keras/keras.json file. Available backends include "tensorflow", "jax", and "torch". An example of setting the backend to JAX:

export KERAS_BACKEND="jax"

In a Jupyter Notebook or Google Colab environment:

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

Note: Ensure to configure the backend prior to importing keras. Once imported, the backend cannot be changed.

Compatibility and Transition

Keras 3 is designed as a drop-in replacement for tf.keras when using the TensorFlow backend. Existing tf.keras code should transition smoothly, requiring only updates to model.save() calls using the .keras format.

For models utilizing custom components (e.g., custom layers or train_step()), converting to a backend-independent implementation usually takes minimal effort. Keras models are versatile in terms of data consumption, supporting tf.data.Dataset pipelines or PyTorch DataLoaders across any backend.

Advantages of Keras 3

• Seamlessly run high-level Keras workflows on any framework, leveraging specific advantages such as JAX's scalability or TensorFlow's production flexibility.
• Develop custom components (like layers and models) applicable across different frameworks.
• Future-proof machine learning code, avoiding dependency on a single framework.
• PyTorch and JAX users gain access to the robust and well-documented features of Keras.

For a comprehensive overview, explore the Keras 3 release announcement.