Build your first neural network

Step 1

The XOR puzzle

XOR means exclusive or: output 1 when two bits differ, otherwise 0. Its four examples are (0,0) → 0, (0,1) → 1, (1,0) → 1, and (1,1) → 0.

Imagine those inputs as the corners of a square. The positive examples sit on opposite corners. No single straight line separates them from the other two corners. A linear model with one output and a threshold cannot solve this pattern.

This is a tiny learning exercise, not a generalization benchmark: we train on all four possible binary inputs. Lesson 6 will measure predictions on genuinely held-out examples.

import torch
from torch import nn

torch.manual_seed(7)
X = torch.tensor([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]])
y = torch.tensor([[0.0], [1.0], [1.0], [0.0]])
print("Inputs:", X.shape, "Targets:", y.shape)

Step 2

Two layers and a bend in the middle

nn.Sequential passes values through its layers in order. nn.Linear(2, 8) makes eight weighted combinations of two inputs, each with its own bias. nn.Tanh() bends those values into the range -1 to 1. nn.Linear(8, 1) combines the hidden values into one final score.

Without a nonlinear activation between them, stacking these linear layers would still be one linear calculation. The nonlinear hidden layer is what allows a more flexible boundary.

The output is a logit, an unrestricted score—not a probability yet. BCEWithLogitsLoss combines a sigmoid transformation with binary cross-entropy in a numerically stable calculation. Do not add a sigmoid before this loss. For this loss, the targets are floating-point 0s and 1s with the same shape as the logits.

An optimizer updates the parameters. Here we use Adam with a learning rate of 0.03. It adjusts update sizes using gradient history; it still depends on gradients from backward().

model = nn.Sequential(
    nn.Linear(2, 8),
    nn.Tanh(),
    nn.Linear(8, 1),
)
loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.03)

print(model)
print("Trainable parameters:", sum(p.numel() for p in model.parameters()))

Step 3

Train, then inspect each prediction

The loop should now look familiar. optimizer.zero_grad() replaces our manual gradient clearing; optimizer.step() replaces the manual subtraction. The model starts in training mode, then switches to evaluation mode for inference.

model.eval() changes the behavior of certain layers, such as dropout; it does not turn gradient tracking off. That is why we also use torch.no_grad(). Our simple Linear/Tanh network has no dropout, but this is a useful habit.

At inference time, sigmoid turns a logit into a value between 0 and 1. We interpret this as the model's estimated probability of class 1 and use a 0.5 threshold. A confident score is not a guarantee that a model is correct.

model.train()
for epoch in range(400):
    optimizer.zero_grad()
    logits = model(X)
    assert logits.shape == y.shape
    loss = loss_fn(logits, y)
    loss.backward()
    optimizer.step()
    if epoch % 100 == 0:
        print(f"Epoch {epoch:3d} | loss before update: {loss.item():.4f}")

model.eval()
with torch.no_grad():
    probabilities = torch.sigmoid(model(X))
    predictions = (probabilities >= 0.5).float()
    accuracy = (predictions == y).float().mean().item()
for inputs, probability, predicted, target in zip(X, probabilities, predictions, y):
    print(inputs.tolist(), f"p(1)={probability.item():.3f}",
          "prediction=", int(predicted.item()), "target=", int(target.item()))
print(f"Training accuracy: {accuracy:.0%}")