Training a Multilayer Perceptron with PyTorch

Welcome back. In the previous lesson, you used autograd to compute gradients and implemented parameter updates with gradient descent and Adam. Now you will connect those pieces into a complete training loop: a multilayer perceptron (MLP) will make predictions, a loss will measure its errors, backpropagation will compute parameter gradients, and an optimizer will update the model.

To keep the geometry visible, the inputs will be individual 3D points rather than images. You will train the MLP to classify whether a point lies inside a unit sphere. This is the last lesson in the foundations module; the next module will examine more carefully how to evaluate a trained model.

A 3D classification task

Write a point as $\mathbf{x}=(x,y,z)$. Its label is $1$ if it lies inside the unit sphere and $0$ otherwise:

We will generate points in a cube surrounding the sphere. The network receives only the three coordinates—not the squared radius used to create the labels. It must learn a useful nonlinear decision boundary from examples.

This is pointwise classification: each point is an independent example. Later in the course, a point-cloud model will receive an entire set of points, where relationships among points and their ordering matter. For now, one point per example makes the training mechanics easier to inspect.

An MLP repeatedly applies affine layers and nonlinear activations. The short video below builds that kind of model and then connects it to a PyTorch training loop. Its examples use digit images; focus on the model and training operations, not the image-specific data preparation.

For our task, the architecture will have widths $3$, $32$, $32$, and $2$, with ReLU between the linear layers. Given a batch of $B$ points, its output has shape $(B,2)$: two logits, or unnormalized class scores, per point. It does not need a final ReLU or softmax.

Loss and the update loop

For a point with true class $y_i$ and logits $z_{i,0},z_{i,1}$, cross-entropy assigns the loss

nn.CrossEntropyLoss calculates the batch mean and handles the log-softmax computation internally. Pass it raw logits, not probabilities produced by a separate softmax. The target must contain integer class indices, here $0$ or $1$.

Read the PyTorch tutorial’s compact explanation of the loss and optimizer before running the code.

Each minibatch follows the same sequence. First the model computes logits and the loss. optimizer.zero_grad() clears gradients left from earlier batches; loss.backward() fills the parameters’ gradient buffers; optimizer.step() uses those gradients to change the parameters. An epoch is one pass through all training examples. Unlike the manual updates in the previous lesson, a PyTorch optimizer performs its update without requiring you to write a torch.no_grad() block around step().

Train the MLP

The following example uses generated data, so it needs only PyTorch—no dataset download or GPU. It creates separate training and held-out points using reproducible random draws. The held-out set gives a quick check on whether learning extends beyond the examples used for updates; a rigorous treatment of data splits comes next.

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

# Generate 3D coordinates in the cube [-1, 1]^3.
data_rng = torch.Generator().manual_seed(17)

def make_points(count):
    x = 2 * torch.rand(count, 3, generator=data_rng) - 1
    y = (x.square().sum(dim=1) < 1).long()
    return x, y

x_train, y_train = make_points(2048)
x_check, y_check = make_points(512)

loader = DataLoader(
    TensorDataset(x_train, y_train),
    batch_size=128,
    shuffle=True,
    generator=torch.Generator().manual_seed(23),
)

torch.manual_seed(7)
model = nn.Sequential(
    nn.Linear(3, 32),
    nn.ReLU(),
    nn.Linear(32, 32),
    nn.ReLU(),
    nn.Linear(32, 2),
)

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=3e-3)

def check_performance():
    model.eval()
    with torch.no_grad():
        logits = model(x_check)
        loss = loss_fn(logits, y_check).item()
        accuracy = (logits.argmax(dim=1) == y_check).float().mean().item()
    return loss, accuracy

print("Input shape:", x_train.shape)
print("Target shape and dtype:", y_train.shape, y_train.dtype)
print("Logit shape:", model(x_train[:8]).shape)
print("Before training (loss, accuracy):", check_performance())

for epoch in range(25):
    model.train()
    loss_total = 0.0

    for batch_x, batch_y in loader:
        optimizer.zero_grad()

        logits = model(batch_x)
        loss = loss_fn(logits, batch_y)

        loss.backward()
        optimizer.step()

        loss_total += loss.item() * len(batch_x)

    if epoch == 0 or (epoch + 1) % 5 == 0:
        mean_loss = loss_total / len(x_train)
        print(f"Epoch {epoch + 1:2d}: training loss = {mean_loss:.4f}")

print("After training (loss, accuracy):", check_performance())

Check the shapes before interpreting the results. x_train has shape $(2048,3)$, and y_train has shape $(2048,)$ with integer dtype. Eight inputs produce logits of shape $(8,2)$. CrossEntropyLoss combines those logits with the corresponding class indices to produce one scalar loss.

The printed training loss should generally fall over the run, while held-out accuracy should improve substantially from its initial value. Exact figures depend on the random draws and PyTorch environment; neither the loss nor accuracy must improve at every individual step. Within an epoch, each batch changes the parameters, so the reported mean combines losses measured at different parameter values.

What the trained model has—and has not—learned

The classifier receives no explicit instruction to calculate $x^2+y^2+z^2$. Its affine layers and ReLUs instead learn a decision surface that approximates the spherical boundary. Errors may persist near that boundary, where a small change in position changes the label.

The target rule depends only on distance from the origin, but an ordinary MLP does not enforce rotation invariance. Its performance on these sampled points therefore does not prove that it will behave consistently under every rotation or on other 3D data. That distinction—between a property of the target and a property guaranteed by the model—will become important in geometric machine learning.

If the code runs but learning stalls, inspect the full chain rather than changing the optimizer immediately: confirm that the inputs and labels align, logits have two columns, targets are integer class indices, the loss is connected to model(batch_x), and backward() occurs before step(). A missing zero_grad() would also cause gradients from successive batches to accumulate unintentionally.

Takeaways

You have trained a small MLP end to end: batches supplied 3D coordinates and labels, the model produced logits, cross-entropy measured their disagreement with the labels, backpropagation computed gradients, and Adam updated the parameters. Monitoring loss and checking predictions on separate examples helped establish that the loop was doing more than merely running without errors.

Next, you will make that check more systematic with train, validation, and test splits—and learn how their metrics reveal overfitting.