From NumPy to PyTorch

2. Tensors and parameters

A tensor is PyTorch’s array. Most of what you know from NumPy works the same way: @, .shape, .sum(0), slicing, broadcasting. Three things are new:

The weights inside nn.Linear are created with requires_grad=True, so they get a .grad after every loss.backward(). Your data does not need one.

Almost everyone makes one mistake the first time: how each side stores a weight matrix. In NumPy you wrote X @ W with W of shape (inputs, outputs). nn.Linear(in, out) stores its weight in the opposite layout, one row per output, and computes x @ W.T + b. For example, nn.Linear(3, 5).weight has shape (5, 3).

3. A model is a class

On the PyTorch side of the lab, the model is one object, model, and you call it like a function: model(X). That object is an instance of a class. A class is a recipe for a box that keeps its own numbers and its own functions. Compare two ways to write the first layer:

def layer(x, W, b):            # a function: you pass the weights in every time
    return x @ W + b

class Layer:                   # a class: the box keeps its weights
    def __init__(self, W, b):  # runs once, when you make the box: layer = Layer(W, b)
        self.W = W             # self means "this box"
        self.b = b
    def forward(self, x):      # the computation, using the box's own weights
        return x @ self.W + self.b

A PyTorch model is a class built the same way, with four rules:

1. __init__ builds the parts: self.l1 = nn.Linear(2, 6). Its first line is always super().__init__() (see the box below).
2. forward(self, x) computes the output from the parts.
3. model(x) calls forward(x) for you. You never call forward by its name.
4. model.parameters() collects every weight of every part you stored as self.something. The optimizer only updates what parameters() returns.

class XorNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Linear(2, 6)
        self.l2 = nn.Linear(6, 1)
    def forward(self, x):
        return self.l2(torch.tanh(self.l1(x)))

model = XorNet()
logits = model(X)              # runs forward(X)

Here is the same idea in NumPy, so it runs in your browser.

Rule 4 has a trap. Level 21 stacks several identical blocks. The natural Python way to keep them is a list: self.blocks = [Block(), Block(), Block()].

4. backward() adds, it doesn’t replace

This one rule causes most beginner bugs. loss.backward() computes the gradients and adds them to whatever is already in .grad. It never clears .grad. Clearing is your job, with opt.zero_grad().

Why add? A parameter can be used in several places, and its total gradient is the sum of the pieces (you build this yourself in level U3). Inside one backward pass, adding is right. Across training steps, it means old gradients stay until you remove them.

Take one weight, w = 1, and one example, x = 2, y = 6, with loss $(w \cdot x - y)^2$. Its gradient is $2(w x - y)\,x = 2 \times (2 - 6) \times 2 = -16$.

Now try it. The buttons do what the PyTorch calls do, on this one weight.

Copy how PyTorch keeps track of gradients, in NumPy. Param below is a tiny class (section 3): a box that holds two numbers, w.data and w.grad (inside the class, self means “this box”). backward adds to w.grad, as PyTorch does. Write zero_grad.

5. Shapes that broadcast without asking

PyTorch broadcasts exactly like NumPy. That is convenient, and it hides a classic bug. The model outputs one score per example as a column, shape (4, 1). Labels are often stored flat, shape (4,).

6. Same numbers, both ways

PyTorch’s gradients are not an approximation. They are exactly the chain rule you wrote by hand. The demo for this level copies the same weights into NumPy and into PyTorch, runs one backward pass in each, and prints both: they match to the last digit. (The demo uses float64 on both sides. With PyTorch’s default float32 they agree to about 7 digits.) Check one of them yourself.

For a 2→3→1 version of the network, with the weights below, PyTorch reports this gradient for the last layer (shown as a column):

Write the NumPy line that produces it.

7. A translation table to keep

When you write PyTorch later, most lines are a NumPy line you already know. This table is what you have used on this page.

Used on this page

for epoch in range(n_epochs):
    for xb, yb in batches(train, size=64):  # shuffled mini-batches (level U4 writes batches)
        logits = model(xb)                   # (64, V): one score per word in the vocabulary
        loss = F.cross_entropy(logits, yb)
        opt.zero_grad()
        loss.backward()
        opt.step()

8. Your first PyTorch run (optional, on your computer)

Optional. Skip this section if you have not installed PyTorch; nothing on this page depends on it, and the level clears without it. It is the only exercise on this page that needs PyTorch (Run it on your computer). Copy these lines into a file first_run.py (or download first_run.py) and run python first_run.py. The seed makes the random starting weights the same on every computer, so everyone gets the same loss.

import torch, torch.nn as nn, torch.nn.functional as F
torch.manual_seed(0)
X = torch.tensor([[0., 0.], [0., 1.], [1., 0.], [1., 1.]])
y = torch.tensor([[0.], [1.], [1.], [0.]])
model = nn.Sequential(nn.Linear(2, 6), nn.Tanh(), nn.Linear(6, 1))
opt = torch.optim.SGD(model.parameters(), lr=0.5)
for step in range(100):
    loss = F.binary_cross_entropy_with_logits(model(X), y)
    opt.zero_grad()
    loss.backward()
    opt.step()
print(round(loss.item(), 3))

9. The loop, by yourself

Last, the skill this level is about: the four steps of training, in the right order. Below, Param is the box from section 4, and the model is a line, w · x + b. backward adds to .grad, like PyTorch. Write the body of the inner loop: three calls, one per line.