The PyTorch backend lets Keras models use PyTorch’s native Module API while retaining Keras layers, losses, optimizers, and custom training steps. section 7.4.2
A custom training step
A PyTorch-backed step runs the model, computes loss, calls loss.backward(), collects weight.value.grad, applies the optimizer without tracking updates, and clears gradients with model.zero_grad(). section 7.4.2
import torch model = get_mnist_model() loss_fn = keras.losses.SparseCategoricalCrossentropy() optimizer = keras.optimizers.Adam() def train_step(inputs, targets): predictions = model(inputs, training=True) #1 loss = loss_fn(targets, predictions) #1 loss.backward() #2 gradients = [weight.value.grad for weight in model.trainable_weights] #3 with torch.no_grad(): #4 optimizer.apply(gradients, model.trainable_weights) #4 model.zero_grad() #5 return loss
Use this MNIST example to see the complete PyTorch-backed training-step sequence, including gradient collection, optimizer application, and gradient clearing. section 7.4.2
Remember these operational rules:
- Backward passes populate gradients on trainable tensors. section 7.4.2
- Clear gradients every step because backward() accumulates values. section 7.4.2
