Keras Documentation

https://keras.io/guides
Admin
This project provides comprehensive documentation and guides for advanced Keras usage, covering...
Tokens:104,159
Snippets:834
Trust Score:10
Update:4 months ago
Show doc for...
Tokens:
Raw
### Setup: Build a Simple Sequential Model

Source: https://keras.io/guides/customizing_quantization

Builds a simple Sequential model for demonstration purposes. This function is used across various quantization examples.

```python
import keras
import numpy as np
from keras import ops

rng = np.random.default_rng()


def get_model():
    """Builds a simple Sequential model for demonstration."""
    return keras.Sequential(
        [
            keras.Input(shape=(10,)),
            keras.layers.Dense(32, activation="relu"),
            keras.layers.Dense(1),
        ]
    )

```

--------------------------------

### Setup for JAX Backend

Source: https://keras.io/guides/custom_train_step_in_jax

Sets the Keras backend to JAX and imports necessary libraries. This setup is required for running the guide.

```python
import os

# This guide can only be run with the JAX backend.
os.environ["KERAS_BACKEND"] = "jax"

import jax
import keras
import numpy as np

```

--------------------------------

### Instantiate Optimizer, Loss Function, and Model

Source: https://keras.io/guides/writing_a_custom_training_loop_in_tensorflow

Prepare the model, optimizer, and loss function for a custom training loop. This setup is essential before starting the training process.

```Python
model = get_model()

# Instantiate an optimizer to train the model.
optimizer = keras.optimizers.Adam(learning_rate=1e-3)
# Instantiate a loss function.
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

```

--------------------------------

### Synchronous Multi-Device Training Setup with JAX

Source: https://keras.io/guides/distributed_training_with_jax

This snippet demonstrates the core setup for synchronous multi-device training in JAX. It includes defining datasets, model, optimizer, loss function, and the gradient computation. It also shows how to initialize model and optimizer states.

```python
# Config
num_epochs = 2
batch_size = 64

train_data, eval_data = get_datasets()
train_data = train_data.batch(batch_size, drop_remainder=True)

model = get_model()
optimizer = keras.optimizers.Adam(1e-3)
loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# Initialize all state with .build()
(one_batch, one_batch_labels) = next(iter(train_data))
model.build(one_batch)
optimizer.build(model.trainable_variables)


# This is the loss function that will be differentiated.
# Keras provides a pure functional forward pass: model.stateless_call
def compute_loss(trainable_variables, non_trainable_variables, x, y):
    y_pred, updated_non_trainable_variables = model.stateless_call(
        trainable_variables, non_trainable_variables, x, training=True
    )
    loss_value = loss(y, y_pred)
    return loss_value, updated_non_trainable_variables


# Function to compute gradients
compute_gradients = jax.value_and_grad(compute_loss, has_aux=True)


# Training step, Keras provides a pure functional optimizer.stateless_apply
@jax.jit
def train_step(train_state, x, y):
    trainable_variables, non_trainable_variables, optimizer_variables = train_state
    (loss_value, non_trainable_variables), grads = compute_gradients(
        trainable_variables, non_trainable_variables, x, y
    )

    trainable_variables, optimizer_variables = optimizer.stateless_apply(
        optimizer_variables, grads, trainable_variables
    )

    return loss_value, (
        trainable_variables,
        non_trainable_variables,
        optimizer_variables,
    )


# Replicate the model and optimizer variable on all devices
def get_replicated_train_state(devices):
    # All variables will be replicated on all devices
    var_mesh = Mesh(devices, axis_names=("_"))
    # In NamedSharding, axes not mentioned are replicated (all axes here)
    var_replication = NamedSharding(var_mesh, P())

    # Apply the distribution settings to the model variables
    trainable_variables = jax.device_put(model.trainable_variables, var_replication)
    non_trainable_variables = jax.device_put(
        model.non_trainable_variables, var_replication
    )
    optimizer_variables = jax.device_put(optimizer.variables, var_replication)

    # Combine all state in a tuple
    return (trainable_variables, non_trainable_variables, optimizer_variables)


num_devices = len(jax.local_devices())
print(f"Running on {num_devices} devices: {jax.local_devices()}")
devices = mesh_utils.create_device_mesh((num_devices,))

# Data will be split along the batch axis
data_mesh = Mesh(devices, axis_names=("batch",))  # naming axes of the mesh
data_sharding = NamedSharding(
    data_mesh,
    P(
        "batch",
    ),  # naming axes of the sharded partition
)

```

--------------------------------

### Setup Keras Backend and Imports

Source: https://keras.io/guides/orbax_checkpoint

Configure the Keras backend to JAX and import necessary libraries. This setup is required for running distributed demos and using JAX with Keras.

```python
import os

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

import shutil

import jax
import keras
import numpy as np

# Simulate 4 CPU devices for the distributed demo.
# Remove this line if using real multi-device hardware.
jax.config.update("jax_num_cpu_devices", 4)

```

--------------------------------

### Setup Imports

Source: https://keras.io/guides/sequential_model

Import necessary modules from Keras for building models.

```python
import keras
from keras import layers
from keras import ops
```

--------------------------------

### Training Loop Setup

Source: https://keras.io/guides/writing_a_custom_training_loop_in_torch

Initializes a Keras model, an optimizer, and a loss function for a training loop. Ensure the model is instantiated before this setup.

```python
# Get a fresh model
model = get_model()

# Instantiate an optimizer to train the model.
optimizer = keras.optimizers.Adam(learning_rate=1e-3)
# Instantiate a loss function.
loss_fn = keras.losses.CategoricalCrossentropy(from_logits=True)
```

--------------------------------

### Setup Imports and Environment

Source: https://keras.io/guides/writing_a_custom_training_loop_in_tensorflow

Imports necessary libraries and sets the Keras backend to TensorFlow. This setup is required for running the TensorFlow-specific code in this guide.

```python
import time
import os

# This guide can only be run with the TensorFlow backend.
os.environ["KERAS_BACKEND"] = "tensorflow"

import tensorflow as tf
import keras
import numpy as np

```

--------------------------------

### Setup imports

Source: https://keras.io/guides/training_with_built_in_methods

Import necessary libraries for Keras, PyTorch, TensorFlow, and numerical operations. These are used for data handling and model building.

```python
# We import torch & TF so as to use torch Dataloaders & tf.data.Datasets.
import torch
import tensorflow as tf

import os
import numpy as np
import keras
from keras import layers
from keras import ops

```

--------------------------------

### Install Keras Nightly

Source: https://keras.io/guides/migrating_to_keras_3

Install the Keras nightly build. This is a prerequisite for using Keras 3.

```bash
!pip install -q keras-nightly

```

--------------------------------

### Quantize and Run Forward Pass Example

Source: https://keras.io/guides/writing_quantization_compatible_layers

Demonstrates quantizing a `SimpleScale` layer to INT8 and performing a forward pass with sample inputs.

```Python
# Sample inputs
rng = np.random.default_rng()
x = rng.random((2, 4)).astype("float32")

layer = SimpleScale()

# Forward pass in float
y_fp = layer(x)


```

--------------------------------

### Install JAX with TPU support

Source: https://keras.io/guides/define_custom_kernel

Installs the latest version of JAX with TPU support and necessary dependencies. Ensure you are running the latest version of libtpu.

```bash
!pip install --upgrade -q "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
```

--------------------------------

### Install Keras and Flax

Source: https://keras.io/guides/keras_nnx_guide

Installs the necessary Keras and Flax libraries. Ensure you are using flax version 0.11.0.

```bash
!pip install -q -U keras
!pip install -q -U flax==0.11.0

```

--------------------------------

### Setup PyTorch Backend and Imports

Source: https://keras.io/guides/writing_a_custom_training_loop_in_torch

Sets the Keras backend to PyTorch and imports necessary libraries including torch, keras, and numpy.

```python
import os

# This guide can only be run with the torch backend.
os.environ["KERAS_BACKEND"] = "torch"

import torch
import keras
import numpy as np
```

--------------------------------

### Setup for PyTorch Backend

Source: https://keras.io/guides/custom_train_step_in_torch

Sets the Keras backend to PyTorch and imports necessary libraries for model building and data handling.

```python
import os

# This guide can only be run with the torch backend.
os.environ["KERAS_BACKEND"] = "torch"

import torch
import keras
from keras import layers
import numpy as np

```

--------------------------------

### Compile Model with Optimizer, Loss, and Metrics

Source: https://keras.io/guides/training_with_built_in_methods

Configure the model's training process by specifying the optimizer, loss function, and metrics to monitor. This setup is required before calling the `fit()` method.

```Python
model.compile(
    optimizer=keras.optimizers.RMSprop(),  # Optimizer
    # Loss function to minimize
    loss=keras.losses.SparseCategoricalCrossentropy(),
    # List of metrics to monitor
    metrics=[keras.metrics.SparseCategoricalAccuracy()]
)
```

--------------------------------

### Setup for JAX Backend

Source: https://keras.io/guides/writing_a_custom_training_loop_in_jax

Sets the Keras backend to JAX and imports necessary libraries including JAX, TensorFlow for data handling, Keras, and NumPy.

```python
import os

# This guide can only be run with the jax backend.
os.environ["KERAS_BACKEND"] = "jax"

import jax

# We import TF so we can use tf.data.
import tensorflow as tf
import keras
import numpy as np

```

--------------------------------

### Quantize, Save, and Load a Model with SimpleScale Layer

Source: https://keras.io/guides/writing_quantization_compatible_layers

Demonstrates how to quantize a sequential model to INT8, save it, and then load the quantized model. This example requires `x` to be defined.

```python
model = keras.Sequential([Input(shape=(4,)), SimpleScale()])
model.build((None, 4))

# Quantize to INT8.
model.quantize("int8")
y_int8 = model(x)
print("SimpleScale INT8 sample:", y_int8[0].numpy())

# Save and load the quantized model.
model.save("simplescale_int8.keras")
loaded = keras.saving.load_model("simplescale_int8.keras")

y_loaded = loaded(x)
print("Loaded INT8 sample:", y_loaded[0].numpy())

```

--------------------------------

### Training Progress Log - Epoch 0

Source: https://keras.io/guides/writing_a_custom_training_loop_in_tensorflow

Example output showing the training loss per batch for the first epoch.

```text
Start of epoch 0
Training loss (for 1 batch) at step 0: 95.3300
Seen so far: 32 samples
Training loss (for 1 batch) at step 100: 2.5622
Seen so far: 3232 samples
Training loss (for 1 batch) at step 200: 3.1138
Seen so far: 6432 samples
Training loss (for 1 batch) at step 300: 0.6748
Seen so far: 9632 samples
Training loss (for 1 batch) at step 400: 1.3308
Seen so far: 12832 samples
Training loss (for 1 batch) at step 500: 1.9813
Seen so far: 16032 samples
Training loss (for 1 batch) at step 600: 0.8640
Seen so far: 19232 samples
Training loss (for 1 batch) at step 700: 1.0696
Seen so far: 22432 samples
Training loss (for 1 batch) at step 800: 0.3662
Seen so far: 25632 samples
Training loss (for 1 batch) at step 900: 0.9556
Seen so far: 28832 samples
Training loss (for 1 batch) at step 1000: 0.7459
Seen so far: 32032 samples
Training loss (for 1 batch) at step 1100: 0.0468
Seen so far: 35232 samples
Training loss (for 1 batch) at step 1200: 0.7392
Seen so far: 38432 samples
Training loss (for 1 batch) at step 1300: 0.8435
Seen so far: 41632 samples
Training loss (for 1 batch) at step 1400: 0.3859
Seen so far: 44832 samples
Training loss (for 1 batch) at step 1500: 0.4156
Seen so far: 48032 samples

```

--------------------------------

### Distributed Training Setup with MirroredStrategy

Source: https://keras.io/guides/distributed_training_with_tensorflow

Initializes a MirroredStrategy to utilize all available GPUs on a single machine. It then defines a model and dataset within the strategy's scope, and proceeds to train and evaluate the model across all replicas.

```python
# Create a MirroredStrategy.
strategy = tf.distribute.MirroredStrategy()
print("Number of devices: {}".format(strategy.num_replicas_in_sync))

# Open a strategy scope.
with strategy.scope():
    # Everything that creates variables should be under the strategy scope.
    # In general this is only model construction & `compile()`.
    model = get_compiled_model()

    # Train the model on all available devices.
    train_dataset, val_dataset, test_dataset = get_dataset()
    model.fit(train_dataset, epochs=2, validation_data=val_dataset)

    # Test the model on all available devices.
    model.evaluate(test_dataset)
```

--------------------------------

### Training Progress Log - Epoch 1

Source: https://keras.io/guides/writing_a_custom_training_loop_in_tensorflow

Example output showing the training loss per batch for the second epoch.

```text
Start of epoch 1
Training loss (for 1 batch) at step 0: 0.4045
Seen so far: 32 samples
Training loss (for 1 batch) at step 100: 0.5983
Seen so far: 3232 samples
Training loss (for 1 batch) at step 200: 0.3154
Seen so far: 6432 samples
Training loss (for 1 batch) at step 300: 0.7911
Seen so far: 9632 samples
Training loss (for 1 batch) at step 400: 0.2607
Seen so far: 12832 samples
Training loss (for 1 batch) at step 500: 0.2303
Seen so far: 16032 samples
Training loss (for 1 batch) at step 600: 0.6048
Seen so far: 19232 samples
Training loss (for 1 batch) at step 700: 0.7041
Seen so far: 22432 samples
Training loss (for 1 batch) at step 800: 0.3669
Seen so far: 25632 samples
Training loss (for 1 batch) at step 900: 0.6389
Seen so far: 28832 samples
Training loss (for 1 batch) at step 1000: 0.7739
Seen so far: 32032 samples
Training loss (for 1 batch) at step 1100: 0.3888
Seen so far: 35232 samples
Training loss (for 1 batch) at step 1200: 0.8133
Seen so far: 38432 samples
Training loss (for 1 batch) at step 1300: 0.2034
Seen so far: 41632 samples
Training loss (for 1 batch) at step 1400: 0.0768
Seen so far: 44832 samples
Training loss (for 1 batch) at step 1500: 0.1544
Seen so far: 48032 samples

```

--------------------------------

### Training Progress Log - Epoch 2

Source: https://keras.io/guides/writing_a_custom_training_loop_in_tensorflow

Example output showing the training loss per batch for the third epoch.

```text
Start of epoch 2
Training loss (for 1 batch) at step 0: 0.1250
Seen so far: 32 samples
Training loss (for 1 batch) at step 100: 0.0152
Seen so far: 3232 samples

```

--------------------------------

### Get Available Devices

Source: https://keras.io/guides/orbax_checkpoint

Prints the number of available devices for distributed computing. This is a preliminary step before setting up distribution.

```python
devices = jax.devices()
print(f"Available devices: {len(devices)}")
```

```text
Available devices: 4

```

--------------------------------

### Learning Rate Scheduling Example

Source: https://keras.io/guides/writing_your_own_callbacks

This snippet demonstrates how a learning rate can be adjusted during training. The learning rate is printed at the beginning of each epoch.

```python
Epoch 10: Learning rate is 0.004999999888241291.
Up to batch 0, the average loss is   15.84.
Up to batch 1, the average loss is   15.76.
Up to batch 2, the average loss is   15.68.
Up to batch 3, the average loss is   15.61.
Up to batch 4, the average loss is   15.53.
Up to batch 5, the average loss is   15.45.
Up to batch 6, the average loss is   15.37.
Up to batch 7, the average loss is   15.29.
Up to batch 8, the average loss is   15.23.
Up to batch 9, the average loss is   15.15.
Up to batch 10, the average loss is   15.08.
Up to batch 11, the average loss is   15.00.
Up to batch 12, the average loss is   14.93.
Up to batch 13, the average loss is   14.86.
Up to batch 14, the average loss is   14.79.
Up to batch 15, the average loss is   14.72.
The average loss for epoch 10 is   14.72 and mean absolute error is    2.37.
```

```python
Epoch 11: Learning rate is 0.004999999888241291.
Up to batch 0, the average loss is   14.65.
Up to batch 1, the average loss is   14.58.
Up to batch 2, the average loss is   14.52.
Up to batch 3, the average loss is   14.45.
Up to batch 4, the average loss is   14.39.
Up to batch 5, the average loss is   14.33.
Up to batch 6, the average loss is   14.26.
Up to batch 7, the average loss is   14.20.
Up to batch 8, the average loss is   14.14.
Up to batch 9, the average loss is   14.08.
Up to batch 10, the average loss is   14.02.
Up to batch 11, the average loss is   13.96.
Up to batch 12, the average loss is   13.90.
Up to batch 13, the average loss is   13.84.
Up to batch 14, the average loss is   13.78.
Up to batch 15, the average loss is   13.72.
The average loss for epoch 11 is   13.72 and mean absolute error is    2.27.
```

```python
Epoch 12: Learning rate is 0.001.
Up to batch 0, the average loss is   13.67.
Up to batch 1, the average loss is   13.60.
Up to batch 2, the average loss is   13.55.
Up to batch 3, the average loss is   13.49.
Up to batch 4, the average loss is   13.44.
Up to batch 5, the average loss is   13.38.
Up to batch 6, the average loss is   13.33.
Up to batch 7, the average loss is   13.28.
Up to batch 8, the average loss is   13.22.
Up to batch 9, the average loss is   13.17.
Up to batch 10, the average loss is   13.12.
Up to batch 11, the average loss is   13.07.
Up to batch 12, the average loss is   13.02.
Up to batch 13, the average loss is   12.97.
Up to batch 14, the average loss is   12.92.
Up to batch 15, the average loss is   12.87.
The average loss for epoch 12 is   12.87 and mean absolute error is    2.19.
```

```python
Epoch 13: Learning rate is 0.0010000000474974513.
Up to batch 0, the average loss is   12.82.
Up to batch 1, the average loss is   12.77.
Up to batch 2, the average loss is   12.72.
Up to batch 3, the average loss is   12.68.
Up to batch 4, the average loss is   12.63.
Up to batch 5, the average loss is   12.58.
Up to batch 6, the average loss is   12.53.
Up to batch 7, the average loss is   12.49.
Up to batch 8, the average loss is   12.45.
Up to batch 9, the average loss is   12.40.
Up to batch 10, the average loss is   12.35.
Up to batch 11, the average loss is   12.30.
Up to batch 12, the average loss is   12.26.
Up to batch 13, the average loss is   12.22.
Up to batch 14, the average loss is   12.17.
Up to batch 15, the average loss is   12.13.
The average loss for epoch 13 is   12.13 and mean absolute error is    2.12.
```

```python
Epoch 14: Learning rate is 0.0010000000474974513.
Up to batch 0, the average loss is   12.09.
Up to batch 1, the average loss is   12.05.
Up to batch 2, the average loss is   12.01.
Up to batch 3, the average loss is   11.97.
Up to batch 4, the average loss is   11.92.
Up to batch 5, the average loss is   11.88.
Up to batch 6, the average loss is   11.84.
Up to batch 7, the average loss is   11.80.
Up to batch 8, the average loss is   11.76.
Up to batch 9, the average loss is   11.72.
Up to batch 10, the average loss is   11.68.
Up to batch 11, the average loss is   11.64.
Up to batch 12, the average loss is   11.60.
Up to batch 13, the average loss is   11.57.
Up to batch 14, the average loss is   11.54.
Up to batch 15, the average loss is   11.50.
The average loss for epoch 14 is   11.50 and mean absolute error is    2.06.
```

--------------------------------

### Setup for Model Serialization

Source: https://keras.io/guides/serialization_and_saving

Import necessary libraries for Keras model operations, including NumPy for data handling and Keras components for model building and operations.

```python
import numpy as np
import keras
from keras import ops
```

--------------------------------

### Layer Serialization and Deserialization

Source: https://keras.io/guides/making_new_layers_and_models_via_subclassing

Demonstrates how to get the configuration of a layer and recreate it from that configuration. This is essential for saving and loading custom layers.

```Python
layer = Linear(64)
config = layer.get_config()
print(config)
new_layer = Linear.from_config(config)
```

--------------------------------

### Setup Imports for Keras Serialization

Source: https://keras.io/guides/customizing_saving_and_serialization

Imports necessary libraries for Keras serialization customization. Ensure these are present before implementing custom saving logic.

```python
import os
import numpy as np
import keras
```

--------------------------------

### Launching Distributed Processes

Source: https://keras.io/guides/distributed_training_with_torch

This snippet demonstrates how to launch multiple Python processes for distributed training using `torch.multiprocessing.start_processes`. It specifies the function to run on each process and the number of processes to start.

```python
if __name__ == "__main__":
    # We use the "fork" method rather than "spawn" to support notebooks
    torch.multiprocessing.start_processes(
        per_device_launch_fn,
        args=(num_gpu,),
        nprocs=num_gpu,
        join=True,
        start_method="fork",
    )

```

--------------------------------

### GAN Training Example

Source: https://keras.io/guides/custom_train_step_in_tensorflow

This snippet shows a typical training step for a Generative Adversarial Network (GAN) using TensorFlow. It demonstrates the output format of a GAN training process.

```python
gan.fit(dataset.take(100), epochs=1)
```

--------------------------------

### Full Model Saving and Loading Example

Source: https://keras.io/guides/serialization_and_saving

Demonstrates saving a compiled Keras model to a `.keras` file and then loading it back. It verifies the reconstruction by comparing predictions from the original and loaded models.

```python
def get_model():
    # Create a simple model.
    inputs = keras.Input(shape=(32,))
    outputs = keras.layers.Dense(1)(inputs)
    model = keras.Model(inputs, outputs)
    model.compile(optimizer=keras.optimizers.Adam(), loss="mean_squared_error")
    return model


model = get_model()

# Train the model.
test_input = np.random.random((128, 32))
test_target = np.random.random((128, 1))
model.fit(test_input, test_target)

# Calling `save('my_model.keras')` creates a zip archive `my_model.keras`.
model.save("my_model.keras")

# It can be used to reconstruct the model identically.
reconstructed_model = keras.models.load_model("my_model.keras")

# Let's check:
np.testing.assert_allclose(
    model.predict(test_input), reconstructed_model.predict(test_input)
)
```

--------------------------------

### Setup Model and Dataset Functions

Source: https://keras.io/guides/distributed_training_with_torch

Defines functions to create a Keras model with batch normalization and dropout, and to load and prepare the MNIST dataset using PyTorch's TensorDataset.

```python
import os

os.environ["KERAS_BACKEND"] = "torch"

import torch
import numpy as np
import keras


def get_model():
    # Make a simple convnet with batch normalization and dropout.
    inputs = keras.Input(shape=(28, 28, 1))
    x = keras.layers.Rescaling(1.0 / 255.0)(inputs)
    x = keras.layers.Conv2D(filters=12, kernel_size=3, padding="same", use_bias=False)(
        x
    )
    x = keras.layers.BatchNormalization(scale=False, center=True)(x)
    x = keras.layers.ReLU()(x)
    x = keras.layers.Conv2D(
        filters=24,
        kernel_size=6,
        use_bias=False,
        strides=2,
    )(x)
    x = keras.layers.BatchNormalization(scale=False, center=True)(x)
    x = keras.layers.ReLU()(x)
    x = keras.layers.Conv2D(
        filters=32,
        kernel_size=6,
        padding="same",
        strides=2,
        name="large_k",
    )(x)
    x = keras.layers.BatchNormalization(scale=False, center=True)(x)
    x = keras.layers.ReLU()(x)
    x = keras.layers.GlobalAveragePooling2D()(x)
    x = keras.layers.Dense(256, activation="relu")(x)
    x = keras.layers.Dropout(0.5)(x)
    outputs = keras.layers.Dense(10)(x)
    model = keras.Model(inputs, outputs)
    return model


def get_dataset():
    # Load the data and split it between train and test sets
    (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

    # Scale images to the [0, 1] range
    x_train = x_train.astype("float32")
    x_test = x_test.astype("float32")
    # Make sure images have shape (28, 28, 1)
    x_train = np.expand_dims(x_train, -1)
    x_test = np.expand_dims(x_test, -1)
    print("x_train shape:", x_train.shape)

    # Create a TensorDataset
    dataset = torch.utils.data.TensorDataset(
        torch.from_numpy(x_train), torch.from_numpy(y_train)
    )
    return dataset

```

--------------------------------

### Custom Layer with Build Configuration Saving/Loading

Source: https://keras.io/guides/customizing_saving_and_serialization

Implement `get_build_config` and `build_from_config` to save and restore custom layer build states, including non-standard arguments like initializers. This example saves and restores a layer initializer.

```python
@keras.saving.register_keras_serializable(package="my_custom_package")
class LayerWithCustomBuild(keras.layers.Layer):
    def __init__(self, units=32, **kwargs):
        super().__init__(**kwargs)
        self.units = units

    def call(self, inputs):
        return keras.ops.matmul(inputs, self.w) + self.b

    def get_config(self):
        return dict(units=self.units, **super().get_config())

    def build(self, input_shape, layer_init):
        # Note the overriding of `build()` to add an extra argument.
        # Therefore, we will need to manually call build with `layer_init` argument
        # before the first execution of `call()`.
        super().build(input_shape)
        self._input_shape = input_shape
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer=layer_init,
            trainable=True,
        )
        self.b = self.add_weight(
            shape=(self.units,),
            initializer=layer_init,
            trainable=True,
        )
        self.layer_init = layer_init

    def get_build_config(self):
        build_config = {
            "layer_init": self.layer_init,
            "input_shape": self._input_shape,
        }  # Stores our initializer for `build()`
        return build_config

    def build_from_config(self, config):
        # Calls `build()` with the parameters at loading time
        self.build(config["input_shape"], config["layer_init"])


custom_layer = LayerWithCustomBuild(units=16)
custom_layer.build(input_shape=(8,), layer_init="random_normal")

model = keras.Sequential(
    [
        custom_layer,
        keras.layers.Dense(1, activation="sigmoid"),
    ]
)

x = np.random.random((16, 8))
y = model(x)

model.save("custom_build_model.keras")
restored_model = keras.models.load_model("custom_build_model.keras")

np.testing.assert_equal(restored_model.layers[0].layer_init, "random_normal")
np.testing.assert_equal(restored_model.built, True)

```

--------------------------------

### Define a Custom Keras Callback

Source: https://keras.io/guides/writing_your_own_callbacks

Create a custom callback by subclassing `keras.callbacks.Callback`. Override methods like `on_train_begin`, `on_epoch_end`, `on_batch_begin`, etc., to hook into different stages of the model's lifecycle. This example logs messages at the start and end of training, epochs, and batches for training, testing, and prediction.

```python
class CustomCallback(keras.callbacks.Callback):
    def on_train_begin(self, logs=None):
        keys = list(logs.keys())
        print("Starting training; got log keys: {}".format(keys))

    def on_train_end(self, logs=None):
        keys = list(logs.keys())
        print("Stop training; got log keys: {}".format(keys))

    def on_epoch_begin(self, epoch, logs=None):
        keys = list(logs.keys())
        print("Start epoch {} of training; got log keys: {}".format(epoch, keys))

    def on_epoch_end(self, epoch, logs=None):
        keys = list(logs.keys())
        print("End epoch {} of training; got log keys: {}".format(epoch, keys))

    def on_test_begin(self, logs=None):
        keys = list(logs.keys())
        print("Start testing; got log keys: {}".format(keys))

    def on_test_end(self, logs=None):
        keys = list(logs.keys())
        print("Stop testing; got log keys: {}".format(keys))

    def on_predict_begin(self, logs=None):
        keys = list(logs.keys())
        print("Start predicting; got log keys: {}".format(keys))

    def on_predict_end(self, logs=None):
        keys = list(logs.keys())
        print("Stop predicting; got log keys: {}".format(keys))

    def on_train_batch_begin(self, batch, logs=None):
        keys = list(logs.keys())
        print("...Training: start of batch {}; got log keys: {}".format(batch, keys))

    def on_train_batch_end(self, batch, logs=None):
        keys = list(logs.keys())
        print("...Training: end of batch {}; got log keys: {}".format(batch, keys))

    def on_test_batch_begin(self, batch, logs=None):
        keys = list(logs.keys())
        print("...Evaluating: start of batch {}; got log keys: {}".format(batch, keys))

    def on_test_batch_end(self, batch, logs=None):
        keys = list(logs.keys())
        print("...Evaluating: end of batch {}; got log keys: {}".format(batch, keys))

    def on_predict_batch_begin(self, batch, logs=None):
        keys = list(logs.keys())
        print("...Predicting: start of batch {}; got log keys: {}".format(batch, keys))

    def on_predict_batch_end(self, batch, logs=None):
        keys = list(logs.keys())
        print("...Predicting: end of batch {}; got log keys: {}".format(batch, keys))
```

--------------------------------

### Install Orbax Checkpointing Library

Source: https://keras.io/guides/orbax_checkpoint

Install the Orbax checkpointing library using pip. This is a prerequisite for using Orbax features in Keras.

```bash
!pip install -q -U orbax-checkpoint

```

--------------------------------

### Instantiate and call a Sequential model

Source: https://keras.io/guides/sequential_model

Demonstrates instantiating a Sequential model and calling it with test data to build the model and create weights. The model is built when it first sees input data.

```Python
model = keras.Sequential(
    [
        layers.Dense(2, activation="relu"),
        layers.Dense(3, activation="relu"),
        layers.Dense(4),
    ]
)  # No weights at this stage!

# At this point, you can't do this:
# model.weights

# You also can't do this:
# model.summary()

# Call the model on a test input
x = ops.ones((1, 4))
y = model(x)
print("Number of weights after calling the model:", len(model.weights))  # 6
```

--------------------------------

### Get Input Shape and Dtype

Source: https://keras.io/guides/functional_api

Access the shape and data type of the input layer.

```python
inputs.shape
```

```python
inputs.dtype
```

--------------------------------

### Prepare and Train a Simple Model

Source: https://keras.io/guides/training_with_built_in_methods

This snippet demonstrates shuffling and batching a dataset, compiling a model, and training it for one epoch.

```python
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(64)

model = get_compiled_model()
model.fit(train_dataset, epochs=1)
```

--------------------------------

### Launch TensorBoard from Command Line

Source: https://keras.io/guides/training_with_built_in_methods

This command launches the TensorBoard visualization tool. Ensure you replace '/full_path_to_your_logs' with the actual directory where your model logs are stored.

```bash
tensorboard --logdir=/full_path_to_your_logs

```

--------------------------------

### Setup TensorFlow for Keras

Source: https://keras.io/guides/distributed_training_with_tensorflow

Sets the Keras backend to TensorFlow and imports necessary libraries. This is a prerequisite for using TensorFlow-specific features.

```python
import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import tensorflow as tf
import keras

```

--------------------------------

### Initialize Optimizer and Run Training Loop

Source: https://keras.io/guides/writing_a_custom_training_loop_in_jax

Build optimizer variables and initialize the training state. Then, iterate through the training dataset, converting TF tensors to NumPy arrays before passing them to the JAX `train_step` function. Log the training loss periodically.

```python
# Build optimizer variables.
optimizer.build(model.trainable_variables)

trainable_variables = model.trainable_variables
non_trainable_variables = model.non_trainable_variables
optimizer_variables = optimizer.variables
state = trainable_variables, non_trainable_variables, optimizer_variables

# Training loop
for step, data in enumerate(train_dataset):
    data = (data[0].numpy(), data[1].numpy())
    loss, state = train_step(state, data)
    # Log every 100 batches.
    if step % 100 == 0:
        print(f"Training loss (for 1 batch) at step {step}: {float(loss):.4f}")
        print(f"Seen so far: {(step + 1) * batch_size} samples")
```

--------------------------------

### ModelParallel and LayoutMap Setup

Source: https://keras.io/guides/distribution

Illustrates setting up `ModelParallel` with a `LayoutMap` for sharding model weights and activations across devices. This is useful for large models that exceed single accelerator memory. It shows how to define tensor layouts using regex-like keys and a `DeviceMesh`.

```python
mesh_2d = keras.distribution.DeviceMesh(
    shape=(2, 4), axis_names=["data", "model"], devices=devices
)
layout_map = keras.distribution.LayoutMap(mesh_2d)
# The rule below means that for any weights that match with d1/kernel, it
# will be sharded with model dimensions (4 devices), same for the d1/bias.
# All other weights will be fully replicated.
layout_map["d1/kernel"] = (None, "model")
layout_map["d1/bias"] = ("model",)

# You can also set the layout for the layer output like
layout_map["d2/output"] = ("data", None)

model_parallel = keras.distribution.ModelParallel(layout_map, batch_dim_name="data")

keras.distribution.set_distribution(model_parallel)

inputs = layers.Input(shape=(28, 28, 1))
y = layers.Flatten()(inputs)
y = layers.Dense(units=200, use_bias=False, activation="relu", name="d1")(y)
y = layers.Dropout(0.4)(y)
y = layers.Dense(units=10, activation="softmax", name="d2")(y)
model = keras.Model(inputs=inputs, outputs=y)

# The data will be sharded across the "data" dimension of the method,

```

--------------------------------

### Execute Training and Resumption

Source: https://keras.io/guides/distributed_training_with_tensorflow

Demonstrates running the training process for one epoch, and then calling it again to resume from the saved checkpoint.

```python
# Running the first time creates the model
run_training(epochs=1)

# Calling the same function again will resume from where we left off
run_training(epochs=1)
```

--------------------------------

### Single-Host, Multi-Device Synchronous Training Setup

Source: https://keras.io/guides/distributed_training_with_torch

This code configures distributed training environment, prepares data loaders, and wraps the model with `DistributedDataParallel`. It's designed to be run across multiple processes, each handling a specific GPU.

```python
# Config
num_gpu = torch.cuda.device_count()
num_epochs = 2
batch_size = 64
print(f"Running on {num_gpu} GPUs")


def setup_device(current_gpu_index, num_gpus):
    # Device setup
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "56492"
    device = torch.device("cuda:{}".format(current_gpu_index))
    torch.distributed.init_process_group(
        backend="nccl",
        init_method="env://",
        world_size=num_gpus,
        rank=current_gpu_index,
    )
    torch.cuda.set_device(device)

def cleanup():
    torch.distributed.destroy_process_group()

def prepare_dataloader(dataset, current_gpu_index, num_gpus, batch_size):
    sampler = torch.utils.data.distributed.DistributedSampler(
        dataset,
        num_replicas=num_gpus,
        rank=current_gpu_index,
        shuffle=False,
    )
    dataloader = torch.utils.data.DataLoader(
        dataset,
        sampler=sampler,
        batch_size=batch_size,
        shuffle=False,
    )
    return dataloader

def per_device_launch_fn(current_gpu_index, num_gpu):
    # Setup the process groups
    setup_device(current_gpu_index, num_gpu)

    dataset = get_dataset()
    model = get_model()

    # prepare the dataloader
    dataloader = prepare_dataloader(dataset, current_gpu_index, num_gpu, batch_size)

    # Instantiate the torch optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    # Instantiate the torch loss function
    loss_fn = torch.nn.CrossEntropyLoss()

    # Put model on device
    model = model.to(current_gpu_index)
    ddp_model = torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[current_gpu_index], output_device=current_gpu_index
    )

    train_model(ddp_model, dataloader, num_epochs, optimizer, loss_fn)

    cleanup()

```

--------------------------------

### Configure and Run AWQ Quantization

Source: https://keras.io/guides/awq_quantization_in_keras

This snippet demonstrates how to set up and apply AWQ quantization to a Keras model using a calibration dataset and specific configuration parameters. Ensure the dataset is representative of your inference domain.

```python
texts = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:1%]")["text"]

calibration_dataset = []
for text in texts:
    for s in text.split("."):
        s = s.strip()
        if s:
            calibration_dataset.append(s + ".")

awq_config = keras.quantizers.AWQConfig(
    dataset=calibration_dataset,
    tokenizer=model.preprocessor.tokenizer,
    weight_bits=4,
    group_size=128,
    num_grid_points=20,
    num_samples=128,
    sequence_length=256,
)

model.quantize("awq", config=awq_config)

outputs = model.generate(prompt, max_length=30)
print(outputs)
```

--------------------------------

### Instantiate Metrics and Prepare Model

Source: https://keras.io/guides/writing_a_custom_training_loop_in_jax

Initializes the model, optimizer, loss function, and metrics for the custom training loop. Metrics like CategoricalAccuracy are instantiated here.

```python
# Get a fresh model
model = get_model()

# Instantiate an optimizer to train the model.
optimizer = keras.optimizers.Adam(learning_rate=1e-3)
# Instantiate a loss function.
loss_fn = keras.losses.CategoricalCrossentropy(from_logits=True)

# Prepare the metrics.
train_acc_metric = keras.metrics.CategoricalAccuracy()
val_acc_metric = keras.metrics.CategoricalAccuracy()
```

--------------------------------

### Quantize Layer to INT8 and Run

Source: https://keras.io/guides/writing_quantization_compatible_layers

Demonstrates quantizing a layer to INT8 and running inference. This is useful for comparing FP32 and INT8 outputs.

```python
layer.quantize("int8")
y_int8 = layer(x)
```

--------------------------------

### Define Input Layer for Image Data

Source: https://keras.io/guides/functional_api

Example of defining an input layer for image data with a shape of (32, 32, 3).

```python
image_input = keras.Input(shape=(32, 32, 3))
```

--------------------------------

### Build Sequential model with predefined input shape

Source: https://keras.io/guides/sequential_model

Illustrates building a Sequential model by first defining an Input layer. This approach ensures the model knows its input shape from the start, allowing for immediate weight creation and summary generation.

```Python
model = keras.Sequential()
model.add(keras.Input(shape=(4,)))
model.add(layers.Dense(2, activation="relu"))

model.summary()
```

--------------------------------

### Eval Loop Example

Source: https://keras.io/guides/writing_a_custom_training_loop_in_jax

This snippet demonstrates an evaluation loop that iterates through a validation dataset, calculates loss, and updates accuracy metrics.

```python
for step, data in enumerate(val_dataset):
    data = (data[0].numpy(), data[1].numpy())
    loss, state = eval_step(state, data)
    # Log every 100 batches.
    if step % 100 == 0:
        print(f"Validation loss (for 1 batch) at step {step}: {float(loss):.4f}")
        _, _, metric_variables = state
        for variable, value in zip(val_acc_metric.variables, metric_variables):
            variable.assign(value)
        print(f"Validation accuracy: {val_acc_metric.result()}")
        print(f"Seen so far: {(step + 1) * batch_size} samples")
```

--------------------------------

### Instantiate Optimizer, Loss Function, and Prepare Dataset

Source: https://keras.io/guides/writing_a_custom_training_loop_in_tensorflow

Sets up the Adam optimizer, Sparse Categorical Crossentropy loss function, and prepares the MNIST dataset for training. Includes shuffling and batching for the training data.

```python
# Instantiate an optimizer.
optimizer = keras.optimizers.Adam(learning_rate=1e-3)
# Instantiate a loss function.
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# Prepare the training dataset.
batch_size = 32
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = np.reshape(x_train, (-1, 784))
x_test = np.reshape(x_test, (-1, 784))

# Reserve 10,000 samples for validation.
x_val = x_train[-10000:]
y_val = y_train[-10000:]
x_train = x_train[:-10000]
y_train = y_train[:-10000]

# Prepare the training dataset.
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)


```

--------------------------------

### Import necessary libraries

Source: https://keras.io/guides/define_custom_kernel

Imports JAX, Keras, and Pallas, setting the Keras backend to JAX. This setup is required for using Pallas kernels.

```python
from functools import partial
import os
import time


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

import jax
from jax.experimental import pallas as pl
import jax.numpy as jnp
import keras
```

--------------------------------

### Dummy Data and Gemma Fine-tuning Step

Source: https://keras.io/guides/keras_nnx_guide

Demonstrates preparing dummy data and performing a single fine-tuning step using the Keras `fit` API with a Gemma language model. This is useful for quick demonstrations or initial testing.

```python
features = np.array(["The quick brown fox jumped.", "I forgot my homework."])
# The model.fit() API works seamlessly!
gemma_lm.fit(x=features, batch_size=2)
print("\n✅ Gemma fine-tuning step completed successfully!")
```

--------------------------------

### Evaluation Output Example

Source: https://keras.io/guides/custom_train_step_in_jax

This shows the typical output format when calling `model.evaluate()` with a custom `test_step`, including metric results and the loss value.

```text
32/32 ━━━━━━━━━━━━━━━━━━━━ 0s 1ms/step - mae: 0.5369 - loss: 0.4170

{'compile_metrics': {'mae': Array(0.5368782, dtype=float32)},
 'loss': Array(0.41702443, dtype=float32)}
```
Tokens:
Raw