Image Created by Author with Playground AI
Before we dive into the details of this exciting development, let’s explore a scenario to understand it better. Picture yourself as a Senior Data Scientist leading a sophisticated image classification project. Your TensorFlow-based model is performing remarkably well. However, as you add more functionalities, you notice that some team members prefer JAX for scalability, while others favor PyTorch for its user-friendly nature. As a team lead, how can you ensure seamless collaboration while maintaining the model’s efficiency across various deep-learning frameworks?
Recognizing this challenge, the Keras team introduces Keras Core—an innovative multi-backend implementation of the Keras API with support for TensorFlow, JAX, and PyTorch. This library will evolve into Keras 3.0 by Fall 2023. But before we jump straight to Keras 3.0, let’s first take a brief look at Keras’ history.
In 2015, François Chollet introduced Keras, an open-source deep-learning library written in Python. This simple yet powerful API quickly gained popularity among researchers, students, and professionals by streamlining complex neural network building. Over time, Keras received significant enhancements, making it even more appealing to the deep-learning community. Eventually, Keras became an integral part of TensorFlow, Google’s cutting-edge deep learning framework. Around the same time, Facebook’s AI Research lab developed PyTorch, known for its intuitive and flexible model building. At the same time, JAX emerged as another powerful framework for high-performance machine learning research. As these frameworks gained momentum, developers started to experience the dilemma of choosing between the frameworks. This led to further fragmentation in the deep learning community.
Acknowledging the challenges faced posed by the fragmented frameworks, the developers of Keras decided to revolutionize the library once again, giving rise to Keras 3.0.
Keras 3.0 empowers you to collaborate with your team effectively. You can develop sophisticated models with a unified front end by combining the strengths of TensorFlow, JAX, and PyTorch. Here are some of the features due to which Keras 3.0 is an absolute game changer:
1. Multi-Backend Support
Keras 3.0 acts as a super-connector, enabling seamless use of TensorFlow, JAX, and PyTorch together. Developers have the freedom to mix and match the best tool for their specific tasks without changing the code.
2. Performance Optimization
Performance Optimization is the key feature of Keras 3.0. By default, Keras 3.0 leverages the XLA (Accelerated Linear Algebra) compilation. XLA Compilation optimizes your mathematical computations making them run much faster on hardware like GPUs and TPUs. It also allows you to dynamically select the best backend for your AI model ensuring optimal efficiency. Such performance optimization features are a big deal allowing you to train more models, experiment more, and get results much quicker.
3. Expanded Ecosystem Surface
Your Keras models can be used as the PyTorch Modules, TensorFlow SavedModels, or part of JAX’s large-scale TPU training infrastructure. This means that you can take advantage of the strengths of each framework. So, with Keras 3.0’s expanded ecosystem you are not locked into a single ecosystem. It is just like a universal adapter that lets you connect your favorite device to any machine.
4. Cross-Framework Low-Level Language
The introduction of the keras_core.ops namespace is a groundbreaking feature that allows you to write custom operations once and use them across different deep-learning frameworks effortlessly. The keras_core.ops provide a set of tools and functions that resemble the popular NumPy API, which is a widely used library for numerical computations in Python. This level of cross-framework compatibility fosters code reusability and encourages collaboration.
5. Progressive Disclosure of Complexity
This design approach of Keras 3.0 sets it apart from other deep learning frameworks. Imagine that you are a beginner and you want to build a simple neural network using Keras 3.0. It presents you with the most straightforward workflows at the start. Once you get comfortable with the basics, you can access all the advanced features and the low-level functionalities. It doesn’t restrict you to predefined workflows only. The beauty of this approach lies in its adaptability and is welcoming for both beginners and seasoned deep learning practitioners.
6. Stateless API for Layers, Models, Metrics, and Optimizers
In the context of deep learning, state refers to the internal variables and parameters that change during the training process. However, JAX operates on the principle of statelessness meaning that the functions do not have mutable variables or internal states. Keras 3.0 embraces the statelessness of JAX via a Stateless API. It allows the essential component of deep learning i.e. layers, models, metrics and optimizers to be designed in a stateless manner. This unique compatibility makes Keras 3.0 an indispensable tool in modern AI development.
Keras Core is compatible with Linux and MacOS systems. Setting up Keras 3.0 is a straightforward process. Here is the step-by-step guide for you to follow:
1. Clone and Navigate to the Repository
Clone the repository into your local system using the following command
git clone https://github.com/keras-team/keras-core.git
Change your root directory to the clone keras-core using:
2. Install Dependencies
Open your interval and run the following command to install the required dependencies.
pip install -r requirements.txt
4. Run the Installation Command
Run the following script to handle the installation process:
python pip_build.py --install
5. Configure the Backend
By default, Keras Core strictly requires TensorFlow as the backend framework but you can configure it using the following two ways:
Option 01: You can set the KERAS_BACKEND environment variable to your preferred backend option.
export KERAS_BACKEND="jax"
Option 02: You can edit the local Keras configuration file located at ~/.keras/keras.json. Open the file in a text editor and change the “backend” option to your preferred backend.
{
"backend": "jax",
"floatx": "float32",
"epsilon": 1e-7,
"image_data_format": "channels_last"
}
6. Verify the Installation
To ensure that Keras Core is correctly installed with your chosen backend, you can test it by importing the library. Open a Python interpreter or a Jupyter Notebook and run the following:
import keras_core as keras
While Keras 3.0 has some limitations, such as its current TensorFlow dependency and limited tf.data support with other backends, the future potential of this framework is promising. Keras has released the beta version as of now and they encourage the developers to provide their valuable feedback. If you are interested to explore more about it, you can find the documentation of Keras Core (Keras 3.0) here. Don’t be afraid to try new ideas. Keras 3.0 is a powerful tool and it’s an exciting time to be part of the evolution.
Kanwal Mehreen is an aspiring software developer with a keen interest in data science and applications of AI in medicine. Kanwal was selected as the Google Generation Scholar 2022 for the APAC region. Kanwal loves to share technical knowledge by writing articles on trending topics, and is passionate about improving the representation of women in tech industry.
