liveBook · Manning

PyTorch backend

As explained in Deep Learning with Python, Third Edition

The execution backend that allows Keras workflows to run on top of PyTorch.

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

How Deep Learning with Python, Third Edition covers it

How deep it goes
Traces forward graphs, backward propagation, gradient storage, optimizer updates, and gradient clearing.
Examples
Uses an MNIST model with SparseCategoricalCrossentropy, Adam, and a one-batch custom train_step.
What it assumes
Assumes familiarity with framework-specific training steps, forward passes, losses, gradients, and optimizers.
What you can do after
Build custom Keras training steps that execute with PyTorch gradients and optimizer updates.
How it connects
Connects Keras models and losses with torch.nn.Module, automatic differentiation, trainable weights, and optimizers.

Related concepts

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with Python, Third Edition. The section links point to the parts of the book it draws on.