Keras 3 Documentation

Use keras.applications.Xception to instantiate the Xception architecture for image classification or feature extraction. The model expects input images of size 299x299 by default.

Basic Usage:

import keras

# Load pre-trained model with ImageNet weights
model = keras.applications.Xception(weights='imagenet')

# Load model without the top classification layers (for transfer learning)
model = keras.applications.Xception(include_top=False, weights='imagenet')

Key Parameters:

• include_top: Set to False to exclude the final fully-connected classification layers (useful for feature extraction or fine-tuning).
• weights: One of None (random initialization), 'imagenet' (pre-trained on ImageNet), or a path to a local weights file.
• input_shape: Optional tuple specifying the input shape (e.g., (150, 150, 3)). Required if include_top=False and the default 299x299 is not desired. Width and height must be no smaller than 71.
• pooling: Optional pooling mode for feature extraction when include_top=False. Options:
  • None: Returns the 4D tensor output of the last convolutional block.
  • 'avg': Applies global average pooling, returning a 2D tensor.
  • 'max': Applies global max pooling, returning a 2D tensor.
• classes: Number of output classes (only used if include_top=True and weights=None). Defaults to 1000.
• classifier_activation: Activation function for the top layer (default 'softmax'). Set to None to return raw logits. When loading pre-trained weights, this can only be None or 'softmax'.

Note on Input Preprocessing: Xception requires specific preprocessing. You must call keras.applications.xception.preprocess_input on your inputs before passing them to the model. This function scales input pixels between -1 and 1.

Example with Preprocessing:

import keras
from keras.applications.xception import preprocess_input

model = keras.applications.Xception(weights='imagenet')

# Assuming 'image' is your input numpy array
preprocessed = preprocess_input(image)
predictions = model.predict(preprocessed)

Sources: keras/src/applications/xception.py

This guide demonstrates how to perform data-parallel distributed training using Keras 3 with the JAX backend on Google Cloud TPU VMs. It covers setting up the environment, building a model, configuring JAX sharding for data and variables, and implementing a custom training loop using Keras' stateless APIs.

Prerequisites

• A Google Cloud TPU VM with at least 8 local devices.
• Follow the official instructions to spin up a TPU VM: https://cloud.google.com/tpu/docs/run-calculation-jax

Setup and Configuration

1. Force the JAX backend by setting the environment variable before importing Keras:

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

2. Import necessary libraries including jax, jax.sharding, and keras.

Model and Data Preparation

• Build a Keras Sequential or functional model (e.g., a CNN for MNIST).
• Load data using tf.data (compatible with Keras) and convert to batches.
• Initialize the model and optimizer state using .build() with a dummy batch:

  model.build(one_batch)
  optimizer.build(model.trainable_variables)

JAX Distribution Setup

• Create a JAX device mesh and sharding configurations:
  • Data Sharding: Split data along the batch axis (P("batch")).
  • Variable Replication: Replicate variables across all devices (P()).
  • Custom Sharding: Optionally shard specific large kernels (e.g., split a Conv2D kernel across 4 devices).

Example mesh and sharding setup:

  from jax.experimental import mesh_utils
  from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

  devices = mesh_utils.create_device_mesh((8,))
  data_mesh = Mesh(devices, axis_names=("batch"))
  data_sharding = NamedSharding(data_mesh, P("batch"))
  
  var_mesh = Mesh(devices, axis_names=("_"))
  var_replication = NamedSharding(var_mesh, P())

Custom Training Loop with Stateless APIs

• Use model.stateless_call for the forward pass and optimizer.stateless_apply for updates. These functions are backend-agnostic and required for functional JAX workflows.
• Define a loss function and a gradient computation function using jax.value_and_grad.
• JIT-compile the training step:

  @jax.jit
  def train_step(train_state, x, y):
      (loss_value, non_trainable_variables), grads = compute_gradients(
          train_state.trainable_variables,
          train_state.non_trainable_variables,
          x,
          y,
      )
      trainable_variables, optimizer_variables = optimizer.stateless_apply(
          train_state.optimizer_variables, grads, train_state.trainable_variables
      )
      return loss_value, TrainingState(
          trainable_variables, non_trainable_variables, optimizer_variables
      )

Inference and State Synchronization

• Run predictions using model.stateless_call with the sharded data.
• After training, update the original Keras model variables using jax.tree_map and variable.assign to ensure the model state is synchronized for subsequent evaluation or saving:

  update = lambda variable, value: variable.assign(value)
  jax.tree_map(update, model.trainable_variables, device_train_state.trainable_variables)
  jax.tree_map(update, model.non_trainable_variables, device_train_state.non_trainable_variables)

• Compile the model and run model.evaluate() to verify the updated state.

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

import jax
import jax.numpy as jnp
import keras
from jax.experimental import mesh_utils
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

# 1. Setup Mesh and Sharding
devices = mesh_utils.create_device_mesh((8,))
data_mesh = Mesh(devices, axis_names=("batch"))
data_sharding = NamedSharding(data_mesh, P("batch"))
var_mesh = Mesh(devices, axis_names=("_"))
var_replication = NamedSharding(var_mesh, P())

# 2. Build and Build State
model = keras.Sequential([...]) # Define model
optimizer = keras.optimizers.Adam(0.01)
model.build(dummy_batch)
optimizer.build(model.trainable_variables)

# 3. Sharding Variables
trainable_variables = [jax.device_put(v, var_replication) for v in model.trainable_variables]
# ... custom sharding logic for specific layers ...

# 4. Stateless Training Step
@jax.jit
def train_step(state, x, y):
    (loss, non_trainable), grads = jax.value_and_grad(
        lambda tv, ntv, x, y: keras.losses.SparseCategoricalCrossentropy()(y, model.stateless_call(tv, ntv, x)[0]),
        has_aux=True
    )(state.trainable_variables, state.non_trainable_variables, x, y)
    tv, opt_vars = optimizer.stateless_apply(state.optimizer_variables, grads, state.trainable_variables)
    return loss, TrainingState(tv, non_trainable, opt_vars)

# 5. Sync State Back to Model
jax.tree_map(lambda v, val: v.assign(val), model.trainable_variables, new_state.trainable_variables)

Sources: examples/demo_jax_distributed.py

To perform single-host, multi-device synchronous distributed training with Keras and JAX, you must manually configure JAX sharding and state management. This guide covers setting up the environment, creating the model and dataset, and implementing a custom training loop that replicates model variables across devices while sharding data.

Prerequisites:

• A machine with multiple GPUs or TPUs (typically 2 to 16).
• JAX backend configured via KERAS_BACKEND=jax environment variable.

Key Steps:

1. Configure Backend: Set os.environ["KERAS_BACKEND"] = "jax" before importing keras.
2. Import Sharding APIs: Import Mesh, NamedSharding, PartitionSpec, and mesh_utils from jax.sharding and jax.experimental.
3. Create Device Mesh: Use mesh_utils.create_device_mesh to define the topology of your devices.
4. Define Sharding Strategies:
   • Variables: Replicate model and optimizer variables across all devices using NamedSharding with an empty PartitionSpec (P()).
   • Data: Shard input data across devices along the batch dimension using NamedSharding with P("batch").
5. Replicate State: Use jax.device_put to distribute the initial model and optimizer state to all devices.
6. Shard Data: Use jax.device_put to shard each batch before passing it to the training step.
7. Custom Training Loop: Implement a loop using model.stateless_call and optimizer.stateless_apply for functional state management.
8. Sync State: After training, manually sync the updated state back to the model variables using variable.assign().

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

import jax
import tensorflow as tf
import keras
from jax.experimental import mesh_utils
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

# 1. Create Device Mesh
devices = jax.local_devices()
num_devices = len(devices)
print(f"Running on {num_devices} devices")
mesh = Mesh(devices, axis_names=("batch",))

# 2. Define Sharding Strategies
# Replicate variables on all devices
var_sharding = NamedSharding(mesh, P())
# Shard data along the batch dimension
data_sharding = NamedSharding(mesh, P("batch"))

# 3. Initialize Model and Optimizer
model = keras.Sequential([...]) # Define your model
optimizer = keras.optimizers.Adam(1e-3)

# Build to initialize variables
dummy_batch = jax.numpy.ones((1, 28, 28, 1))
model.build(dummy_batch)
optimizer.build(model.trainable_variables)

# 4. Replicate State to Devices
trainable_vars = jax.device_put(model.trainable_variables, var_sharding)
non_trainable_vars = jax.device_put(model.non_trainable_variables, var_sharding)
optimizer_vars = jax.device_put(optimizer.variables, var_sharding)

train_state = (trainable_vars, non_trainable_vars, optimizer_vars)

# 5. Training Loop (Simplified)
@jax.jit
def train_step(state, x, y):
    trainable, non_trainable, opt_vars = state
    # Functional forward pass
    y_pred, updated_non_trainable = model.stateless_call(
        trainable, non_trainable, x
    )
    loss = keras.losses.sparse_categorical_crossentropy(y, y_pred)
    
    # Compute gradients
    (loss, updated_non_trainable), grads = jax.value_and_grad(
        lambda t, nt, x, y: keras.losses.sparse_categorical_crossentropy(
            y, model.stateless_call(t, nt, x)[0]
        ), has_aux=True
    )(trainable, non_trainable, x, y)
    
    # Functional optimizer step
    trainable, opt_vars = optimizer.stateless_apply(opt_vars, grads, trainable)
    
    return (trainable, updated_non_trainable, opt_vars), loss

# Process batch
x_sharded = jax.device_put(x_batch.numpy(), data_sharding)
y_sharded = jax.device_put(y_batch.numpy(), data_sharding)
train_state, loss = train_step(train_state, x_sharded, y_sharded)

# 6. Sync State Back to Model
trainable_vars, non_trainable_vars, optimizer_vars = train_state
for var, val in zip(model.trainable_variables, trainable_vars):
    var.assign(val)
for var, val in zip(model.non_trainable_variables, non_trainable_vars):
    var.assign(val)

Sources: guides/distributed_training_with_jax.py

To perform single-host, multi-device synchronous training with Keras and TensorFlow, use the tf.distribute.MirroredStrategy API. This setup replicates the model across multiple GPUs on a single machine, keeping replicas synchronized after each batch to ensure convergence behavior matches single-device training.

Steps to configure:

1. Set the KERAS_BACKEND environment variable to tensorflow.
2. Instantiate a MirroredStrategy object.
3. Open a strategy.scope() context.
4. Create and compile your Keras model inside the scope. This ensures variables are mirrored correctly.
5. Call fit() and evaluate() inside the scope (or ensure the first call to fit() creates variables within the scope).
6. Use tf.data.Dataset objects for data loading to optimize performance.

Example:

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

import tensorflow as tf
import keras

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

# Open a strategy scope
with strategy.scope():
    # Create and compile the model inside the scope
    inputs = keras.Input(shape=(784,))
    x = keras.layers.Dense(256, activation="relu")(inputs)
    outputs = keras.layers.Dense(10)(x)
    model = keras.Model(inputs, outputs)
    model.compile(
        optimizer=keras.optimizers.Adam(),
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=[keras.metrics.SparseCategoricalAccuracy()],
    )

    # Prepare data using tf.data.Dataset
    (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
    x_train = x_train.reshape(-1, 784).astype("float32") / 255
    x_test = x_test.reshape(-1, 784).astype("float32") / 255
    
    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(32)
    val_dataset = tf.data.Dataset.from_tensor_slices((x_train[-10000:], y_train[-10000:])).batch(32)
    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)

    # Train the model
    model.fit(train_dataset, epochs=2, validation_data=val_dataset)
    model.evaluate(test_dataset)

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

import tensorflow as tf
import keras

strategy = tf.distribute.MirroredStrategy()
print('Number of devices: {}'.format(strategy.num_replicas_in_sync))

with strategy.scope():
    inputs = keras.Input(shape=(784,))
    x = keras.layers.Dense(256, activation="relu")(inputs)
    outputs = keras.layers.Dense(10)(x)
    model = keras.Model(inputs, outputs)
    model.compile(
        optimizer=keras.optimizers.Adam(),
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=[keras.metrics.SparseCategoricalAccuracy()],
    )

    train_dataset = tf.data.Dataset.from_tensor_slices((...)).batch(32)
    model.fit(train_dataset, epochs=2)

Sources: guides/distributed_training_with_tensorflow.py

To support sample_weight and class_weight arguments in fit(), unpack them from the data argument in train_step() and pass them to compute_loss() and metric.update_state().

Implementation:

1. Check the length of data. If it has 3 elements, unpack x, y, sample_weight. Otherwise, set sample_weight = None.
2. Pass sample_weight to self.compute_loss().
3. Pass sample_weight to metric.update_state() for non-loss metrics.

Example:

class CustomModel(keras.Model):
    def train_step(self, data):
        if len(data) == 3:
            x, y, sample_weight = data
        else:
            sample_weight = None
            x, y = data

        self.zero_grad()
        y_pred = self(x, training=True)
        loss = self.compute_loss(
            y=y,
            y_pred=y_pred,
            sample_weight=sample_weight,
        )
        loss.backward()

        trainable_weights = [v for v in self.trainable_weights]
        gradients = [v.value.grad for v in trainable_weights]

        with torch.no_grad():
            self.optimizer.apply(gradients, trainable_weights)

        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(loss)
            else:
                metric.update_state(y, y_pred, sample_weight=sample_weight)

        return {m.name: m.result() for m in self.metrics}

model = CustomModel(inputs, outputs)
model.compile(optimizer="adam", loss="mse", metrics=["mae"])
model.fit(x, y, sample_weight=sw, epochs=3)

def train_step(self, data):
    if len(data) == 3:
        x, y, sample_weight = data
    else:
        sample_weight = None
        x, y = data

    self.zero_grad()
    y_pred = self(x, training=True)
    loss = self.compute_loss(
        y=y,
        y_pred=y_pred,
        sample_weight=sample_weight,
    )
    loss.backward()

    trainable_weights = [v for v in self.trainable_weights]
    gradients = [v.value.grad for v in trainable_weights]

    with torch.no_grad():
        self.optimizer.apply(gradients, trainable_weights)

    for metric in self.metrics:
        if metric.name == "loss":
            metric.update_state(loss)
        else:
            metric.update_state(y, y_pred, sample_weight=sample_weight)

    return {m.name: m.result() for m in self.metrics}

Sources: guides/custom_train_step_in_torch.py

To customize the behavior of model.evaluate(), override the test_step(self, data) method in your keras.Model subclass. The implementation is similar to train_step() but uses training=False for the forward pass.

Steps:

1. Unpack data into x and y.
2. Perform a forward pass: y_pred = self(x, training=False).
3. Compute the loss using self.compute_loss().
4. Update metrics using metric.update_state().
5. Return a dictionary mapping metric names to their current values.

Example:

class CustomModel(keras.Model):
    def test_step(self, data):
        x, y = data
        y_pred = self(x, training=False)
        loss = self.compute_loss(y=y, y_pred=y_pred)
        
        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(loss)
            else:
                metric.update_state(y, y_pred)
        
        return {m.name: m.result() for m in self.metrics}

# Usage
model = CustomModel(inputs, outputs)
model.compile(loss="mse", metrics=["mae"])
model.evaluate(x, y)

Sources: unknown