Title: What Is Keras Core? - PyImageSearch

URL Source: https://pyimagesearch.com/2023/07/24/what-is-keras-core/

Published Time: 2023-07-24T13:00:00+00:00

Markdown Content:
* * *

**Table of Contents**
---------------------

*   [What Is Keras Core?](#h2BPTitle)

*   [Configuring Your Development Environment](#h3DevelopmentEnvironment)
*   [Let’s Talk about Keras!](#h3Talk)
*   [Going Beyond with Keras Core](#h3Beyond)
*   [The Power of Keras Core: Expanding Your Deep Learning Horizons](#h3Power)
*   [Show Me Some Code](#h3ShowMeSomeCode)
*   [JAX Harnessing `model.fit()`](#h3JAXHarnessing)

*   [Imports and Setup](#h4ImportsAndSetup)
*   [Data Pipeline](#h4DataPipeline)
*   [Build a Custom Model](#h4BuildCM)
*   [Build the Image Classification Model](#h4BuildICM)
*   [Train the Model](#h4TrainModel)
*   [Evaluation](#h4Evaluation)

*   [Summary](#h2Summary)

*   [References](#h3References)
*   [Citation Information](#h3Citation)

* * *

In this tutorial, you will learn about [Keras Core](https://github.com/keras-team/keras-core), the new Keras Team repository that allows switching from the Keras backend to [TensorFlow](https://www.tensorflow.org/), [PyTorch](https://pytorch.org/), and [JAX](https://jax.readthedocs.io/en/latest/).

[![Image 1](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20393'%3E%3C/svg%3E)](https://pyimagesearch.com/wp-content/uploads/2023/07/keras-core-featured.png)

This lesson introduces Keras Core and its multi-backend functionalities.

**To learn how to navigate the Keras Core library,** **_just keep reading._**

* * *

[**What Is Keras Core?**](#TOC)
-------------------------------

In short, Keras Core is the first effort to end the 100-year war of Deep Learning Frameworks and unite forces for a better and “open” artificial intelligence (AI) future.

We are obviously joking, but our opening statement is not completely untrue. To quote the actual definition from the [Keras Core repository](https://github.com/keras-team/keras-core/tree/main):

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

This means you can switch your backends like Iron Man switching his suits as and when he needs them. But the only difference is, you are not limited to just Iron Man suits, but also Batman, Green Lantern, and Blue Beetle … you get the point.

If you are already sold and need to see some cold hard proof (code examples), jump to the next section.

However, read on if you are interested in recapping Keras and understanding how `keras_core` came to be.

* * *

### [**Configuring Your Development Environment**](#TOC)

To follow this guide, you need to have the `keras-core` library installed on your system.

Luckily, `keras-core` is pip-installable:

$ pip install keras-core

* * *

[**Let’s Talk about Keras!**](#TOC)
-----------------------------------

First, let’s clarify what [**Keras**](https://keras.io/) is. Keras is a _user-friendly_ tool written in Python for Deep Learning. It’s designed to be used with [**TensorFlow**](https://www.tensorflow.org/), another major player in the AI field.

Think of Keras as your personal assistant in the realm of machine learning. Its job is to make your life as a developer _easier_ and more _efficient_. Here are some reasons why Keras is so cool:

*   **It’s User-friendly:** Keras strives to reduce your workload, allowing you to focus on the bigger picture. It’s designed to be easy to use, quick for debugging, elegant and concise in coding, maintainable, and easily deployable.
*   **It’s Adaptable:** Keras believes in starting simple and **gradually unveiling complexity**. The idea is to make basic tasks quick and easy and allow for more complex tasks to be accomplished step by step.
*   **It’s Powerful:** Keras delivers professional-grade performance and scalability. Impressive, right? Big names like NASA, YouTube, and Waymo trust it. This means your YouTube video suggestions and some of the most sophisticated self-driving vehicles rely on Keras.

In short, Keras’s blend of _simplicity_ and _extensibility_ makes it an attractive starting point for everyone, from AI novices to seasoned deep learning engineers.

* * *

[**Going Beyond with Keras Core**](#TOC)
----------------------------------------

But, it would be amazing if we could do more. How about creating models that can train on any data pipeline and run on any platform? Enter Keras Core!

Keras Core is a new and improved version of Keras. It supports multiple platforms, including [TensorFlow](https://www.tensorflow.org/), [JAX](https://jax.readthedocs.io/en/latest/), and [Py](https://pytorch.org/)[T](https://pytorch.org/)[orch](https://pytorch.org/). It’s intended to be a **drop-in** replacement for `tf.keras` when using the TensorFlow backend.

If you have a `tf.keras` model that doesn’t include custom components, you can effortlessly switch it to run on JAX or PyTorch. Even if it does include custom components, you can usually tweak it to work with any backend in just a few minutes.

Plus, Keras Core models can digest datasets in any format, no matter what backend you’re using. You can train your models with existing [`tf.data.Dataset`](https://www.tensorflow.org/api_docs/python/tf/data/Dataset) pipelines or [Torch DataLoaders](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html).

So, Keras is evolving beyond a conventional deep learning framework. It’s becoming a code-based User Interface where you can toggle between different backends, features, and components to make your work easier. Imagine having a LEGO set of machine learning tools — Keras Core lets you construct massive, complex systems from simple, interchangeable components.

* * *

[**The Power of Keras Core: Expanding Your Deep Learning Horizons**](#TOC)
--------------------------------------------------------------------------

**Seamless Framework Switching:** With Keras Core, you’re no longer tied to one framework. You can effortlessly shift your high-level Keras workflows between TensorFlow, JAX, PyTorch, and others. Fancy the scalability of JAX or the production-ready features of TensorFlow? With Keras Core, you’re free to enjoy the best of both worlds.

**Custom Components Across Frameworks:** Need a custom layer, model, or metric? Keras Core lets you integrate custom components that work across all frameworks, aiding in developing unique solutions for your deep learning challenges.

**Universal Training Loops:** Keras Core allows you to train a Keras model using a training loop crafted from scratch in TensorFlow, JAX, or PyTorch. This flexibility ensures your model is trained most efficiently for your particular project.

**Integration with Native Models:** Keras Core extends flexibility by letting you incorporate a Keras model as part of a PyTorch-native Module or a JAX-native model function. This means you can smoothly integrate the advantages of Keras with your existing PyTorch or JAX workflows.

**Future-Proof Code:** With Keras Core, you can make your machine learning code future-proof, evading the risk of being locked into a single framework.

**Boosting User Experience:** PyTorch users, rejoice! Keras Core lets you tap into the powerful and user-friendly nature of Keras. JAX users, you’re not left out either! Keras Core provides you with a fully equipped, battle-hardened, and well-documented modeling and training library.

In essence, Keras Core elevates your deep learning experience, delivering a versatile and flexible tool geared for your success.

* * *

[**Show Me Some Code**](#TOC)
-----------------------------

Wow! That is some neat engineering work done by the Keras team. Let’s dive deeper into the package, see how to switch backends, and take full advantage of the `keras_core`.

Let’s start with installing the package.

Keras Core is compatible with the Linux and MacOS systems. To install a local development version:

1.  `pip install -r requirements.txt`
2.  `python pip_build.py --install`

Note that Keras Core strictly requires TensorFlow, particularly because it uses `tf.nest` to handle nested Python structures. In the future, the repository will most likely make all backend frameworks optional.

Now that you have the package installed, let’s talk about the **backend** first. To configure your backend, export the environment variable `KERAS_BACKEND` or edit your local config file at `~/.keras/keras.json`. Available backend options are: “tensorflow”, “jax”, and “torch”. Example:

export KERAS\_BACKEND="jax"

Now we know that we can switch backends with the `KERAS_BACKEND` environment variable. So, what is stopping us from changing the backend inside a Python script?

import os
os.environ\["KERAS\_BACKEND"\] = "jax"

Neat, isn’t it?

Now for the fun part. If you are a PyImageSearch reader, you will remember the blog post series on JAX.

*   [Learning JAX in 2023: Part 1 — The Ultimate Guide to Accelerating Numerical Computation and Machine Learning](https://pyimagesearch.com/2023/02/20/learning-jax-in-2023-part-1-the-ultimate-guide-to-accelerating-numerical-computation-and-machine-learning/)
*   [Learning JAX in 2023: Part 2 — JAX’s Power Tools `grad`, `jit`, `vmap`, and `pmap`](https://pyimagesearch.com/2023/02/27/learning-jax-in-2023-part-2-jaxs-power-tools-grad-jit-vmap-and-pmap/)
*   [Learning JAX in 2023: Part 3 — A Step-by-Step Guide to Training Your First Machine Learning Model with JAX](https://pyimagesearch.com/2023/04/17/learning-jax-in-2023-part-3-a-step-by-step-guide-to-training-your-first-machine-learning-model-with-jax/)

We also covered an Image Classification model in JAX, where we talked about writing a custom train step and building the model using [Flax](https://flax.readthedocs.io/en/latest/).

*   [JAX Image Classification LiveStream](https://youtube.com/live/zM308DRgjPA)

Let us address the elephant in the room. JAX is good because it is fast. It allows you to run your experiments in jet speed (jit compilation), but writing a custom train step and using another library (Flax) for modeling is like embracing an unorganized airport with no signs (poor design and documentation of Flax) before boarding your jet.

You could skip all of that and actually go to the runway and board the jet with your familiar car.

In Deep Learning terms, this means having the ability to use the Keras modeling and training APIs while using JAX functions for speedup and gradient calculations!

* * *

[**JAX Harnessing `model.fit()`**](#TOC)
----------------------------------------

**A note of credit:** The official [guide](https://github.com/keras-team/keras-core/blob/main/guides/writing_a_custom_training_loop_in_jax.py) at Keras Core repository heavily inspired the following example.

* * *

### [**Imports and Setup**](#TOC)

import os
os.environ\["KERAS\_BACKEND"\] = "jax"
import keras\_core

import jax
import tensorflow as tf
import tensorflow\_datasets as tfds
from matplotlib import pyplot as plt

We set the `KERAS_BACKEND` environment variable to “jax”. This is done for us to use JAX as the backend for `keras_core`.

The other imports are necessary for this example.

* * *

### [**Data Pipeline**](#TOC)

train\_ds, val\_ds, test\_ds = tfds.load(
    "cifar10",
    split=\["train\[:90%\]", "train\[90%:\]", "test"\]
)

This line loads the [CIFAR-10](https://www.tensorflow.org/datasets/catalog/cifar10) dataset and splits it into 90% for training (`train_ds`), 10% for validation (`val_ds`), and the original test set for testing (`test_ds`).

def map\_fn(inputs):
    image = tf.cast(inputs\["image"\], "float32") / 255.
    label = inputs\["label"\]

    return image, label

This function, `map_fn`, is used to preprocess the inputs from the CIFAR-10 dataset:

*   `image = tf.cast(inputs["image"], "float32") / 255.`: It takes the “image” field from the inputs, casts it to `float32` data type, and normalizes the pixel values to the range `[0, 1]` by dividing by `255`.
*   `label = inputs["label"]`: It takes the “label” field from the inputs, which represents the class of the image.

The function returns a tuple of the processed image and its corresponding label.

train\_ds = (
    train\_ds
    .map(map\_fn)
    .shuffle(32 \* 10)
    .batch(32)
    .prefetch(tf.data.AUTOTUNE)
)

val\_ds = (
    val\_ds
    .map(map\_fn)
    .shuffle(32 \* 10)
    .batch(32)
    .prefetch(tf.data.AUTOTUNE)
)

test\_ds = (
    test\_ds
    .map(map\_fn)
    .batch(32)
    .prefetch(tf.data.AUTOTUNE)
)

These lines of code apply transformations to the training, validation, and test datasets:

*   `.map(map_fn)`: Applies the `map_fn` function to each element in the dataset. This function scales the images and extracts the labels.
*   `.shuffle(32 * 10)`: Shuffles the dataset. This is important for the training dataset to ensure the model gets data randomly. The number `32 * 10` is the buffer size for the shuffle operation. Notice how we do not shuffle the `test_ds`.
*   `.batch(32)`: Batches the dataset into groups of 32. This means that the model will update its weights after seeing 32 samples.
*   `.prefetch(tf.data.AUTOTUNE)`: Prefetches the data for faster consumption. It allows the dataset to asynchronously fetch batches while the model is training.
*   `tf.data.AUTOTUNE` allows TensorFlow to automatically choose the number of batches to prefetch based on available resources, optimizing for throughput.

The training and validation datasets are shuffled, but the test dataset is not because shuffling has no effect during model evaluation.

* * *

### [**Build a Custom Model**](#TOC)

In a JAX workflow, the training and testing steps are designed to be **stateless**, meaning they don’t maintain any internal state. Instead, the state is passed explicitly as a parameter to these functions. This approach differs from traditional workflows where the state (like model weights or optimizer state) is maintained internally and updated in place (stateful).

The state typically includes the trainable parameters of the model, non-trainable parameters, and the optimizer state. It is updated after each step and passed back. This explicit handling of state is beneficial for parallel and distributed computing where maintaining and synchronizing internal state can be challenging.

In the training step, the state is used to compute the model’s predictions, calculate the loss, and update the model parameters using backpropagation. The updated state is then returned.

In the testing step, the state is used to compute the model’s predictions on the test data. The state remains unchanged during testing as we don’t perform any updates to the model parameters.

This stateless design, combined with JAX’s functional programming approach, provides a flexible and efficient way to implement complex machine learning workflows.

class CustomModel(keras\_core.Model):
    def compute\_loss\_and\_updates(
        self,
        trainable\_variables,
        non\_trainable\_variables,
        x,
        y,
        training=False,
    ):
        y\_pred, non\_trainable\_variables = self.stateless\_call(
            trainable\_variables,
            non\_trainable\_variables,
            x,
            training=training,
        )
        loss = self.compute\_loss(x, y, y\_pred)
        return loss, (y\_pred, non\_trainable\_variables)

The above code snippet defines a custom model class `CustomModel` that inherits from `keras_core.Model`.

The `compute_loss_and_updates` method is a custom method that calculates the loss and updates the model’s variables. Here’s what it does:

*   `y_pred, non_trainable_variables = self.stateless_call(trainable_variables, non_trainable_variables, x, training=training)`: This line computes the model’s predictions (`y_pred`) given the inputs (`x`) and the current state of trainable and non-trainable variables. The `stateless_call` method is a feature of Keras Core that allows for stateless computation of the model’s forward pass.
*   `loss = self.compute_loss(x, y, y_pred)`: This line computes the loss using a method `compute_loss` that you would need to define in your `CustomModel` class. It takes the true labels (`y`), the model’s predictions (`y_pred`), and the inputs (`x`).
*   `return loss, (y_pred, non_trainable_variables)`: The method returns the computed loss and a tuple containing the model’s predictions and the updated non-trainable variables.

This setup allows for a clear separation of the forward pass, loss computation, and the updates to the model’s variables, providing a flexible framework for customizing the training process.

    def train\_step(self, state, data):
        # Unpack the current state
        (
            trainable\_variables,
            non\_trainable\_variables,
            optimizer\_variables,
            metrics\_variables,
        ) = state

        # Unpack the data
        x, y = data

        # Get the gradient function.
        grad\_fn = jax.value\_and\_grad(
            self.compute\_loss\_and\_updates, has\_aux=True
        )

        # Compute the gradients.
        (loss, (y\_pred, non\_trainable\_variables)), grads = grad\_fn(
            trainable\_variables,
            non\_trainable\_variables,
            x,
            y,
            training=True,
        )

        # Update trainable variables and optimizer variables.
        (
            trainable\_variables,
            optimizer\_variables,
        ) = self.optimizer.stateless\_apply(
            optimizer\_variables, grads, trainable\_variables
        )

        # Update metrics.
        new\_metrics\_vars = \[\]
        for metric in self.metrics:
            this\_metric\_vars = metrics\_variables\[
                len(new\_metrics\_vars) : len(new\_metrics\_vars)
                + len(metric.variables)
            \]
            if metric.name == "loss":
                this\_metric\_vars = metric.stateless\_update\_state(
                    this\_metric\_vars, loss
                )
            else:
                this\_metric\_vars = metric.stateless\_update\_state(
                    this\_metric\_vars, y, y\_pred
                )
            logs = metric.stateless\_result(this\_metric\_vars)
            new\_metrics\_vars += this\_metric\_vars

        # Return metric logs and updated state variables.
        state = (
            trainable\_variables,
            non\_trainable\_variables,
            optimizer\_variables,
            new\_metrics\_vars,
        )
        return logs, state

This code defines the `train_step` method for the `CustomModel` class. This method is responsible for performing a **single training step**, including _forward pass_, _loss computation_, _backpropagation_, and _updating the model parameters and metrics_. Here’s a breakdown:

*   The current state of the model and the data for this step are passed as arguments. The state includes trainable, non-trainable, optimizer, and metrics variables.
*   The `jax.value_and_grad` function is used to create a function (`grad_fn`) that computes both the value of the `compute_loss_and_updates` function and its gradients with respect to the trainable variables.
*   The gradients are computed by calling `grad_fn` with the current trainable and non-trainable variables and the data. The function also returns the loss and the updated non-trainable variables.
*   The `stateless_apply` method of the optimizer is then used to update the trainable and optimizer variables using the computed gradients.
*   The metrics are updated by iterating over each metric in the model’s metrics. For each metric, its state variables are updated using its `stateless_update_state` method, and the result is computed using its `stateless_result` method.
*   Finally, the method returns the logs (containing the metric results) and the updated state.

This method encapsulates a single step of the training process in a stateless manner, making it compatible with JAX’s functional programming model.

    def test\_step(self, state, data):
        # Unpack the data.
        x, y = data
        (
            trainable\_variables,
            non\_trainable\_variables,
            metrics\_variables,
        ) = state

        # Compute predictions and loss.
        y\_pred, non\_trainable\_variables = self.stateless\_call(
            trainable\_variables,
            non\_trainable\_variables,
            x,
            training=False,
        )
        loss = self.compute\_loss(x, y, y\_pred)

        # Update metrics.
        new\_metrics\_vars = \[\]
        for metric in self.metrics:
            this\_metric\_vars = metrics\_variables\[
                len(new\_metrics\_vars) : len(new\_metrics\_vars)
                + len(metric.variables)
            \]
            if metric.name == "loss":
                this\_metric\_vars = metric.stateless\_update\_state(
                    this\_metric\_vars, loss
                )
            else:
                this\_metric\_vars = metric.stateless\_update\_state(
                    this\_metric\_vars, y, y\_pred
                )
            logs = metric.stateless\_result(this\_metric\_vars)
            new\_metrics\_vars += this\_metric\_vars

        # Return metric logs and updated state variables.
        state = (
            trainable\_variables,
            non\_trainable\_variables,
            new\_metrics\_vars,
        )
        return logs, state

This code defines the `test_step` method for the `CustomModel` class. This method is responsible for performing a single step of evaluation or testing. Here’s a breakdown:

*   The current state of the model and the data for this step are passed as arguments. The state includes trainable variables, non-trainable variables, and metrics variables.
*   The `stateless_call` method of the model is used to compute the predictions and update the non-trainable variables.
*   The loss is computed using the `compute_loss` method of the model.
*   The metrics are updated by iterating over each metric in the model’s metrics. For each metric, its state variables are updated using its `stateless_update_state` method, and the result is computed using its `stateless_result` method.
*   Finally, the method returns the logs (containing the metric results) and the updated state.

This method encapsulates a single step of the testing process in a stateless manner, making it compatible with JAX’s functional programming model.

* * *

### [**Build the Image Classification Model**](#TOC)

conv2d\_kwargs = {
    "kernel\_size": (3, 3),
    "activation": "relu",
    "padding": "same",
}

This line creates a dictionary `conv2d_kwargs` with parameters for a 2D convolutional layer. It specifies a 3×3 kernel size, ReLU activation, and the same padding to maintain input dimensions.

inputs = keras\_core.Input(shape=(32, 32, 3), name="input\_layer")
x = inputs

for filters in \[32, 64, 128\]:
    x = keras\_core.layers.Conv2D(filters=filters, \*\*conv2d\_kwargs)(x)
    x = keras\_core.layers.BatchNormalization()(x)
    x = keras\_core.layers.Conv2D(filters=filters, strides=(2, 2), \*\*conv2d\_kwargs)(x)
    x = keras\_core.layers.BatchNormalization()(x)
    x = keras\_core.layers.Dropout(0.25)(x)

x = keras\_core.layers.GlobalAveragePooling2D()(x)
x = keras\_core.layers.Dense(128, activation="relu")(x)
x = keras\_core.layers.Dropout(0.25)(x)
outputs = keras\_core.layers.Dense(10, activation="softmax", name="output\_layer")(x)

Here, we construct a Convolutional Neural Network (CNN) model using Keras Core.

*   It starts by defining an input layer that accepts images of shape `(32, 32, 3)`.
*   Then, it creates three blocks of layers, each consisting of two convolutional layers followed by batch normalization and dropout. The number of filters in the convolutional layers increases from `32` to `64` to `128` across the blocks.
*   After the blocks, it applies global average pooling to the feature maps, followed by a dense layer with ReLU activation and another dropout layer.
*   Finally, it adds an output layer with 10 units (for 10 classes) and softmax activation for multi-class classification.

model = CustomModel(inputs, outputs, name="image\_classification\_model")
model.compile(
    optimizer=keras\_core.optimizers.Adam(learning\_rate=1e-4),
    loss="sparse\_categorical\_crossentropy",
    metrics=\["accuracy"\]
    )
model.summary()

Finally, we instantiate our custom model using the Functional API of Keras. We then compile the model with the Adam optimizer, sparse categorical cross-entropy as the loss function, and accuracy as the metric for evaluation. The model’s architecture is then displayed with the `model.summary()` (figure below) method.

[![Image 2](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%200%200'%3E%3C/svg%3E)](https://lh6.googleusercontent.com/i4UXmFfsK5L2qgJRs5xGawIQGvpRVM_Gtt-hM9WghmKfEsuYefxeUWiEmsduQ4Ceic5RKbNYL3oRunCgazUxOaqor8Wqb1T8gqk6meYRynQy1GHQcoVZk-p67IrosicSn4Ju9bwUDh_6sLlr7Hhq5d8)

Source: Image by the Authors

* * *

### [**Train the Model**](#TOC)

history = model.fit(
    train\_ds,
    epochs=20,
    validation\_data=val\_ds
)

And lastly, for the moment we have been waiting for, we call the legendary `model.fit()` on the JAX workflow. Below is a figure showing that our model is indeed training.

[![Image 3](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%200%200'%3E%3C/svg%3E)](https://lh5.googleusercontent.com/typpw6ulemj3aJU9HhmfVOA7CsfY5AA2QkMyosyI8n84_zyGOnwYRNTd48-FbWZITUix9Ub0cDq-Ve2Bid-3UXW_xKAZWnz9-Xj76DC-ci1MUXRDhcltC1nBn4v6FJj_ABTHuP2En4NtrEg12vyDupM)

Source: Image by the Authors

plt.plot(history.history\["accuracy"\], label="train acc")
plt.plot(history.history\["val\_accuracy"\], label="val acc")
plt.legend()
plt.title("Accuracy Plot")
plt.show()

We generate the plot to visualize the model’s training and validation accuracy (figure below) across the epochs.

[![Image 4](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%200%200'%3E%3C/svg%3E)](https://lh6.googleusercontent.com/ZX98ciVpoQw614rJkTAsXuuFrvmsfw4yoDbxnDzrhk0rWvpMh0wwjHCbvnKjopEoMhIW6B5E_PORTAywAYPXG2uF5wXhioVU6gA-7fDxZoV4bmKIv-3pyZPi6F1Amjitc0uHclEIDFkIvzCJcSawfss)

Source: Image by the Authors

plt.plot(history.history\["loss"\], label="train loss")
plt.plot(history.history\["val\_loss"\], label="val loss")
plt.legend()
plt.title("Loss Plot")
plt.show()

We generate another plot to visualize the model’s training and validation loss (figure below) across the epochs.

[![Image 5](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%200%200'%3E%3C/svg%3E)](https://lh4.googleusercontent.com/8sbAZ3TH38sYXk6pAPGOkk3TScf6pjKz8Dn2DfhYxHqUxgVA_fMNag8iRChiXkkjJthLKoaH8xNym6AtDNETd0mcmnSz1G3DqD65vphjKHuxsqDMkWZzVJZ8uPHFz_yO8-YNc5WlHDCqT3LRLXwNEZY)

Source: Image by the Authors

* * *

### [**Evaluation**](#TOC)

loss, accuracy = model.evaluate(test\_ds)
print(f"{loss=:0.4f}")
print(f"{accuracy=:0.4f}")

We evaluate the model on the test dataset to measure its loss and accuracy (screenshot below in the figure), finally assessing the model’s performance.

[![Image 6](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%2066'%3E%3C/svg%3E)](https://lh4.googleusercontent.com/lQpjdlszbvqRiZnvofLssiZyqW-nRKooSJgDgQdjZJCZKXwvSOtse6tEI2uKzJPmQDYTPp5k6mYyn8nvmzGXKL3BWVFYzrqJgev3cApkNwTWyOJGHs4x-wkoyeBSZZbSnffVoS-uRLj8dQvKFir7UP0)

Source: Image by the Authors

The project output is visualized in the following gif (shown in the figure below). We can see that there are a few misclassifications, but a few classes have been predicted accurately. A pretty good feat for a model trained for only **20 epochs**.

[![Image 7](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20512%20384'%3E%3C/svg%3E)](https://pyimagesearch.com/wp-content/uploads/2023/07/classification.gif)

Source: Image by the Authors

* * *

### What's next? We recommend [PyImageSearch University](https://pyimagesearch.com/pyimagesearch-university/?utm_source=blogPost&utm_medium=bottomBanner&utm_campaign=What%27s%20next%3F%20I%20recommend).

**Course information:**  
84 total classes • 114+ hours of on-demand code walkthrough videos • Last updated: February 2024  
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

**I strongly believe that if you had the right teacher you could _master_ computer vision and deep learning.**

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s _not_ the case.

All you need to master computer vision and deep learning is for someone to explain things to you in _simple, intuitive_ terms. _And that’s exactly what I do_. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to _successfully_ and _confidently_ apply computer vision to your work, research, and projects. Join me in computer vision mastery.

**Inside PyImageSearch University you'll find:**

*   ✓ **84 courses** on essential computer vision, deep learning, and OpenCV topics
*   ✓ **84 Certificates** of Completion
*   ✓ **114+ hours** of on-demand video
*   ✓ **Brand new courses released _regularly_**, ensuring you can keep up with state-of-the-art techniques
*   ✓ **Pre-configured Jupyter Notebooks in Google Colab**
*   ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
*   ✓ Access to **centralized code repos for _all_ 536+ tutorials** on PyImageSearch
*   ✓ **Easy one-click downloads** for code, datasets, pre-trained models, etc.
*   ✓ **Access** on mobile, laptop, desktop, etc.

[Click here to join PyImageSearch University](https://pyimagesearch.com/pyimagesearch-university/?utm_source=blogPost&utm_medium=bottomBanner&utm_campaign=What%27s%20next%3F%20I%20recommend)

* * *

[**Summary**](#TOC)
-------------------

Alright, let’s wrap this up! In this tutorial, we unpacked a small section of things possible using Keras Core.

*   **Keras Core:** We started by introducing Keras Core, the deep learning library that’s as adaptable as a chameleon. It can switch between TensorFlow, PyTorch, and JAX backends faster than you can say “neural network”!
*   **JAX:** We then picked JAX as our backend of choice. Why? Because it’s like the cool new kid on the block who’s great at number crunching and loves complex computations.
*   **Data Loading:** Next, we juggled with the CIFAR-10 dataset, splitting it into training, validation, and test sets. We also normalized and batched the data, because, let’s face it, nobody likes to train on messy data.
*   **Building the Model:** We then played architect and built our own custom Keras model.
*   **JAX Workflow:** We then danced with JAX workflows, integrating them into our Keras Core model. We integrated the traditional `model.fit` API with a stateless train and test step, making our model compatible with JAX’s `value_and_grad` function.
*   **Model Training:** After compiling our model and setting it up with the right optimizer, loss function, and metrics, we trained our model. It worked! Losses went down, accuracies went up, and Keras is now officially compatible with other frameworks like JAX.
*   **Model Evaluation:** Finally, we put our model to the test, literally. We evaluated it on the test set to see if all that training paid off. Spoiler alert: it did!
*   **Visualizing Progress:** And, of course, what’s a journey without some selfies along the way? We captured our model’s progress over time with some cool plots of accuracy and loss.

But with Keras Core, this is just the tip of the iceberg (sorry, Twitter thread makers). The full range of capabilities now includes:

*   **Mix and Match:** With Keras Core, you can use any framework you choose. You can enjoy the speed and scalability of JAX or the production-ready features of TensorFlow. It’s like having your cake and eating it too!
*   **Custom Creations:** You can write your own components, like layers, models, or metrics. Then, you can use these in any framework. It’s like making and using your own Lego blocks in any Lego set.
*   **Training Flexibility:** You can train a Keras model in any framework. Whether you’re a TensorFlow, JAX, or PyTorch fan, Keras Core has got you covered.
*   **Integration Ease:** You can take a Keras model and integrate it into a PyTorch or JAX model. It’s like adding a new player to your favorite sports team.
*   **Future-Proofing:** With Keras Core, you’re not locked into one framework. This means your code is ready for whatever the future of machine learning brings.
*   **For PyTorch Users:** If you’re a PyTorch user, rejoice! You can now enjoy the power and user-friendliness of Keras.
*   **For JAX Users:** If you’re a JAX user, good news! You now have access to a robust, well-documented library for modeling and training (as you just saw).

In short, Keras Core is like a universal adapter for machine learning frameworks. It allows you to choose, customize, and integrate, making your machine learning journey smoother and more flexible.

But it is also much more than that. Keras Core is the first step a big framework has shown in joining hands and acknowledging that each framework has its strengths and weaknesses and playing as a team will let us play longer.

* * *

### [**References**](#TOC)

*   [GitHub – keras-team/keras-core: A multi-backend implementation of the Keras API, with support for TensorFlow, JAX, and PyTorch](https://github.com/keras-team/keras-core)
*   [Keras Core: Keras for TensorFlow, JAX, and PyTorch](https://keras.io/keras_core/)

* * *

### [**Citation Information**](#TOC)

**A. R. Gosthipaty and R. Raha.** “What Is Keras Core?” _PyImageSearch_, P. Chugh, S. Huot, K. Kidriavsteva, and A. Thanki, eds., 2023, [https://pyimg.co/bm6vo](https://pyimg.co/bm6vo)

@incollection{ARG-RR\_2023\_WhatIsKerasCore,
  author = {Aritra Roy Gosthipaty and Ritwik Raha},
  title = {What Is Keras Core?},
  booktitle = {PyImageSearch},
  editor = {Puneet Chugh and Susan Huot and Kseniia Kidriavsteva and Abhishek Thanki},
  year = {2023},
  url = {https://pyimg.co/bm6vo},
}

* * *

![Image 8: Featured Image](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20500%20276'%3E%3C/svg%3E)

### Unleash the potential of computer vision with Roboflow - Free!

*   Step into the realm of the future by [signing up or logging into your Roboflow account](https://roboflow.com/?ref=pyimagesearch). Unlock a wealth of innovative dataset libraries and revolutionize your computer vision operations.
*   Jumpstart your journey by choosing from our broad array of datasets, or benefit from PyimageSearch’s comprehensive library, crafted to cater to a wide range of requirements.
*   Transfer your data to Roboflow in any of the 40+ compatible formats. Leverage cutting-edge model architectures for training, and deploy seamlessly across diverse platforms, including API, NVIDIA, browser, iOS, and beyond. Integrate our platform effortlessly with your applications or your favorite third-party tools.
*   Equip yourself with the ability to train a potent computer vision model in a mere afternoon. With a few images, you can import data from any source via API, annotate images using our superior cloud-hosted tool, kickstart model training with a single click, and deploy the model via a hosted API endpoint. Tailor your process by opting for a code-centric approach, leveraging our intuitive, cloud-based UI, or combining both to fit your unique needs.
*   Embark on your journey today with absolutely no credit card required. Step into the future with Roboflow.

[Join Roboflow Now](https://roboflow.com/?ref=pyimagesearch)

* * *

**To download the source code to this post (and be notified when future tutorials are published here on PyImageSearch), _simply enter your email address in the form below!_**

![Image 9](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20410%20567'%3E%3C/svg%3E)

#### Download the Source Code and FREE 17-page Resource Guide

Enter your email address below to get a .zip of the code and a **FREE 17-page Resource Guide on Computer Vision, OpenCV, and Deep Learning.** Inside you'll find my hand-picked tutorials, books, courses, and libraries to help you master CV and DL!