Özer Cem KELAHMETKELAHMET
Tools · Deep learning

Keras

Keras is a high-level API for building neural networks quickly. Keras 3 runs on JAX, PyTorch or TensorFlow backends. Learn its design and when to use it.

Overview

Keras is a high-level API for deep learning. It lets you describe a neural network as a stack or graph of layers, choose a loss and an optimizer, and train with a single call to fit. The goal has always been to make common work quick and readable, so that a beginner can build a working model in a few lines and an expert can move fast on standard architectures.

Keras 3 is a multi-backend library. The same model code can run on JAX, PyTorch or TensorFlow, and you choose the engine with a setting. That means a team can prototype on one backend, train on whichever is fastest for their hardware and export to the runtime their deployment needs, without rewriting the model.

Keras is like the automatic gearbox in a car. You still decide where to go and how fast, but you do not have to manage every gear change. If you want manual control for a tricky hill, you can still take it, by writing custom layers or a custom training step.

What it is

Keras is made of a few core building blocks.

Building block Purpose
Layers Reusable units such as Dense, Conv2D, LSTM, MultiHeadAttention, LayerNormalization
Sequential API A plain stack of layers for simple models
Functional API Models as graphs with several inputs, outputs or shared layers
Model subclassing Writing a model as a Python class for full flexibility
compile, fit, evaluate, predict The built-in training and inference workflow
Losses, metrics, optimizers Ready-made components, including Adam and AdamW
Callbacks Early stopping, checkpoints, learning-rate schedules, logging
keras.ops A backend-neutral NumPy-style API plus neural network functions
KerasHub Pretrained models and components such as BERT, Whisper and Stable Diffusion

The key idea in Keras 3 is that every built-in layer, loss and optimizer is written against keras.ops. Because keras.ops is implemented for each backend, the whole library works unchanged on all of them. Custom code that also sticks to keras.ops gains the same portability.

Where it came from

François Chollet released Keras on 27 March 2015, as part of a research project called ONEIROS. It is licensed under Apache 2.0. Early Keras was itself multi-backend: up to version 2.3 it could run on Theano, TensorFlow, Microsoft's CNTK and PlaidML.

As TensorFlow grew, Keras became its official high-level API, available as tf.keras, and TensorFlow 2.0 in 2019 made it the centre of model building. From version 2.4 onwards, standalone Keras supported only TensorFlow.

The next turn came with Keras 3, a full rewrite released in late 2023 after a public beta of about five months. It restored multi-backend support, this time for JAX, TensorFlow and PyTorch. Later releases added OpenVINO as an inference-only backend in version 3.8. From TensorFlow 2.16 onwards, Keras 3 is the default Keras that TensorFlow installs, and the old Keras 2 lives on as the tf_keras package in maintenance mode. Keras 3 releases continued through 2026.

How it works

Choosing a backend

You select the backend with the KERAS_BACKEND environment variable, or in the ~/.keras/keras.json file, before importing Keras. Options are "jax", "tensorflow", "torch" and, for inference, "openvino".

Building a model

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

model = keras.Sequential([
    keras.Input(shape=(20,)),
    layers.Dense(64, activation="relu"),
    layers.Dropout(0.2),
    layers.Dense(1, activation="sigmoid"),
])
model.compile(optimizer="adam", loss="binary_crossentropy", metrics=[keras.metrics.AUC()])

Training

model.fit(x, y, epochs=10, validation_split=0.2) runs the loop. Behind the scenes Keras calls the backend to compute the forward pass, the loss, the gradients and the optimizer update. On JAX it compiles the step with XLA; on TensorFlow it uses tf.function; on PyTorch it uses PyTorch's autograd. fit accepts NumPy arrays, Pandas DataFrames, tf.data datasets and PyTorch DataLoader objects on any backend.

Customising

When the standard loop is not enough, you can override train_step in a model, write a custom layer with build and call methods, or write the whole loop in the backend's native style. Keras 3 also exposes stateless versions of its objects, such as layer.stateless_call, so they fit inside pure JAX functions.

Saving and exporting

Models save to the .keras format, which holds no backend-specific content, so a model saved under one backend can be reloaded under another. A Keras model can also be used as a PyTorch module, exported as a TensorFlow SavedModel for serving, or turned into a stateless JAX function.

Key concepts and terms

Term Meaning
Backend The engine that executes Keras code: JAX, TensorFlow, PyTorch or OpenVINO
Layer A unit that holds weights and transforms inputs
Sequential model A linear stack of layers
Functional API Defining models as graphs of layer calls
compile Attaching optimizer, loss and metrics to a model
fit The built-in training loop
Callback An object that runs at points during training
keras.ops Backend-neutral operations used by all Keras components
.keras file The portable saved-model format
KerasHub Library of pretrained models and presets
tf_keras The legacy Keras 2 package kept for older code

A worked example

This scenario is illustrative. A subscription business wants to predict which customers will cancel in the next month. It has 120,000 customer records with 20 numeric features, and about 8 per cent of customers cancelled.

  1. A data scientist scales the features and splits the data 80 to 20 for training and validation.
  2. She builds the small Sequential model shown above: one hidden layer of 64 units, dropout of 0.2 and a sigmoid output.
  3. Because the classes are unbalanced, she passes class_weight to fit so mistakes on the cancelling customers count more.
  4. She trains with the JAX backend, using an EarlyStopping callback that watches validation AUC and stops after three epochs without improvement.
  5. She compares against a gradient boosting baseline. On tabular data like this the boosted trees may well be competitive, so she keeps whichever performs better on the validation set.
  6. If the neural model wins, she saves it as a .keras file and the engineering team reloads it under the TensorFlow backend to export a SavedModel for their existing serving stack.

The model code is written once; only the backend setting changes between training and export.

Strengths and limitations

Strengths:

  • Very readable code and quick prototyping.
  • The same model runs on JAX, TensorFlow or PyTorch.
  • Sensible defaults for losses, metrics and training loops.
  • A large catalogue of layers and pretrained models through KerasHub.
  • Portable .keras save format across backends.

Limitations:

  • Highly unusual research ideas may need backend-specific code.
  • Performance varies between backends for the same model, so benchmarking is needed.
  • Some features, such as the keras.distribution API, have been more complete on JAX than on other backends.
  • Keras 2 and Keras 3 cannot be mixed in one program, which complicates migrations.

How it compares with alternatives

Tool Level Backends Best for
Keras High-level API JAX, TensorFlow, PyTorch Fast building of standard models
PyTorch Low to mid-level framework Itself Research and full control
PyTorch Lightning Training structure over PyTorch PyTorch Organised PyTorch at scale
JAX with Flax Functional, low to mid-level JAX High-performance research
scikit-learn Classical ML API NumPy Non-deep tabular models

Practical guidance and best practices

  1. Set KERAS_BACKEND at the very top of your program, before importing Keras.
  2. Start with the Sequential or Functional API; subclass only when the graph style gets in the way.
  3. Write custom layers with keras.ops, not backend functions, if you want them portable.
  4. Always pass keras.Input with a shape so models build early and show clear summaries.
  5. Use callbacks for early stopping and checkpoints rather than manual loops.
  6. Benchmark the same model on more than one backend before committing to one for training.
  7. Save in the .keras format and test reloading before deploying.
  8. Use KerasHub presets for transfer learning instead of building large models from scratch.

Common mistakes

  • Importing Keras before setting the backend, so the default backend is used silently.
  • Mixing tf.keras from Keras 2 with Keras 3 objects in the same model.
  • Writing custom layers with TensorFlow or PyTorch calls and then expecting them to work on another backend.
  • Forgetting to scale inputs, which slows or destabilises training.
  • Reading training accuracy rather than validation metrics when judging the model.

Frequently asked questions

Does Keras only work with TensorFlow?

No. Keras 3 runs on JAX, TensorFlow and PyTorch, with OpenVINO available for inference. Keras was tied to TensorFlow only between versions 2.4 and 3.0.

Which backend should I use?

It depends on your model and hardware. The Keras team reports that JAX is often fastest, but not always. Pick the backend that matches your deployment stack, then benchmark if training speed matters.

Can I use a Keras model inside PyTorch code?

Yes. With the PyTorch backend, a Keras model behaves as a PyTorch module, so it can be used in a PyTorch training loop or combined with other PyTorch modules.

Is Keras suitable for research?

For many projects, yes, especially with custom layers and custom train_step methods. For very new ideas that need fine control over every operation, researchers often prefer the underlying frameworks directly.

What is KerasHub?

It is the library of pretrained Keras models and building blocks, with presets for language, vision and audio models. It grew out of the earlier KerasNLP library and gathers pretrained work under one name.

How do I migrate from Keras 2?

Most tf.keras code runs on Keras 3 with the TensorFlow backend after small changes, and models should be re-saved in the .keras format. Code that must stay on Keras 2 can use the tf_keras package.

Glossary

  • Activation: a function such as ReLU or sigmoid applied after a layer.
  • Callback: code that runs at points during training.
  • Epoch: one full pass over the training data.
  • Functional API: building models by calling layers on symbolic inputs.
  • Optimizer: the algorithm that updates weights, such as Adam.
  • Preset: a named pretrained configuration in KerasHub.
  • Stateless call: running a layer with weights passed in explicitly, without side effects.
  • XLA: the compiler JAX and TensorFlow use to speed up Keras training.

Summary

Keras, created by François Chollet in 2015, is a high-level API that makes building and training neural networks simple and readable. Keras 3, released in late 2023, runs the same models on JAX, TensorFlow or PyTorch, with a portable save format and a backend-neutral operations layer. It suits fast development of standard models and teams that want freedom over runtime, while unusual research may still call for the underlying frameworks.

It belongs to Frameworks in my stack, under Deep learning. Others in the same group:

Type to search.
    ↑↓ navigate↵ open