PyTorch: the zero-grad, forward, backward and step sequence

pythontorchindependent-practice

Saved on this device

Full lesson · 4 learning steps Open any exercise for focused practice

Step 1 of 4 · intro card

PyTorch: the zero-grad, forward, backward and step sequence

Concept and contract

Observe a real SGD parameter update with an analytic one-example target.

A basic training update has distinct stages: clear the intended gradient window, compute predictions and scalar loss, backpropagate derivatives, then ask an optimizer to update registered parameters.

Calling backward alone does not change weights. Keeping those stages explicit helps distinguish a learning-rate problem, missing registration, stale gradients and an update that was never executed.

Mechanism and worked example

The helper initializes a one-weight linear model to a known value, predicts one from one input and compares it with target two using mean squared error..

The helper initializes a one-weight linear model to a known value, predicts one from one input and compares it with target two using mean squared error.

SGD with learning rate one tenth moves a zero weight to 0.4. The observed returned weight is selected from the real parameter after optimizer.step, not computed as a separate hand-written formula.

Complete worked example

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def one_update(initial, learning_rate):
    model = nn.Linear(1, 1, bias=False)
    with torch.no_grad():
        model.weight.fill_(initial)
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
    x = torch.tensor([[1.0]])
    target = torch.tensor([[2.0]])
    optimizer.zero_grad(set_to_none=True)
    loss = nn.MSELoss()(model(x), target)
    loss.backward()
    optimizer.step()
    return model.weight.item()

Reasoning: The helper initializes a one-weight linear model to a known value, predicts one from one input and compares it with target two using mean squared error.

Counterexamples and common mistakes

The counterexample omits optimizer.step, leaving the initial weight unchanged despite a valid backward call..

The counterexample omits optimizer.step, leaving the initial weight unchanged despite a valid backward call.

Another common error is resetting gradients after backward but before step, which removes the information the optimizer needs. Do not place forward loss computation under no_grad during training; the initialization block uses it only for a deliberate parameter assignment.

Counterexample: a runnable framework mistake

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def one_update(initial, learning_rate):
    model = nn.Linear(1, 1, bias=False)
    with torch.no_grad():
        model.weight.fill_(initial)
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
    x = torch.tensor([[1.0]])
    target = torch.tensor([[2.0]])
    optimizer.zero_grad(set_to_none=True)
    loss = nn.MSELoss()(model(x), target)
    loss.backward()
    pass  # no update applied
    return model.weight.item()

Reasoning: The counterexample omits optimizer.step, leaving the initial weight unchanged despite a valid backward call.

Independent practice: one bounded framework change

Add an optional steps argument and repeat the complete update sequence that many times, clearing gradients before every iteration.

Add an optional steps argument and repeat the complete update sequence that many times, clearing gradients before every iteration. Preserve the known initialization and fixed data.

With two steps from zero at learning rate one tenth, the weight should reach 0.72. The task adds a small loop rather than a full epoch system, checkpoint scheduler or distributed training layer.

Step 1: Repeat the complete SGD update with a gradient reset per step. Step 2: Wrap zero_grad through step in one loop.

Practice starter: solve before reading the solution

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def one_update(initial, learning_rate):
    # Remove this placeholder, then adapt the retained baseline:
    # Repeat the complete SGD update with a gradient reset per step.
    raise NotImplementedError
    model = nn.Linear(1, 1, bias=False)
    with torch.no_grad():
        model.weight.fill_(initial)
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
    x = torch.tensor([[1.0]])
    target = torch.tensor([[2.0]])
    optimizer.zero_grad(set_to_none=True)
    loss = nn.MSELoss()(model(x), target)
    loss.backward()
    optimizer.step()
    return model.weight.item()

Reasoning: Repeat the complete SGD update with a gradient reset per step. Wrap zero_grad through step in one loop.

Hint 1: Keep model initialization outside the loop.
Hint 2: Check two-step weight 0.72 from zero.

Solution, tests and trade-offs

Baseline cases compare positive, zero and nonzero-initialization updates.

Baseline cases compare positive, zero and nonzero-initialization updates. Practice checks one and two steps and zero iterations.

Finite tolerances account for float32 representation. These analytic values make a missing step or unintended accumulation visible, but they do not establish convergence for an arbitrary architecture or justify a learning rate on a real dataset.

Practice solution and acceptance checks

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def one_update(initial, learning_rate, steps=1):
    model = nn.Linear(1, 1, bias=False)
    with torch.no_grad():
        model.weight.fill_(initial)
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
    x = torch.tensor([[1.0]])
    target = torch.tensor([[2.0]])
    for _ in range(steps):
        optimizer.zero_grad(set_to_none=True)
        loss = nn.MSELoss()(model(x), target)
        loss.backward()
        optimizer.step()
    return model.weight.item()

Reasoning: Baseline cases compare positive, zero and nonzero-initialization updates. Practice checks one and two steps and zero iterations.

Interview reasoning and next step

When training larger models, track losses and validation metrics over time, inspect gradient norms when necessary and preserve reproducible configuration..

When training larger models, track losses and validation metrics over time, inspect gradient norms when necessary and preserve reproducible configuration.

Optimizer state matters when resuming momentum-based methods, while this simple SGD fixture has a minimal state footprint. Keep evaluation separate from updates, and do not interpret a falling training loss as proof that the model generalizes.

Optional stretch — after the core task works: Add a bounded loss history and compare training with held-out evaluation.

Computer practice: https://applydjinn.com/typorithm/python-frameworks-lab-20261105.zip. Install the pinned dependencies in an isolated environment, run verify.py --track torch, then edit the retained exercise baseline and check it with --practice pt-optimizer. The lab uses owned in-memory databases, local framework clients and tiny CPU tensors. It never reads the working ApplyDjinn database, contacts production feeds or downloads models.

Observe a real SGD parameter update with an analytic one-example target.