PyTorch basics

Overview

Imagine we give a model the number 2 and want it to answer 5. At first it answers 2.

We score how wrong that is, work out which way to adjust the model’s numbers, and adjust them a little.

Its next answer is 3.5: closer, not yet right.

That is one learning update. Repeating it over many examples is training.

We’ll take three stops: organize the numbers, build an adjustable model, then follow the update.

Who does what:

• PyTorch does the arithmetic and works out the adjustments.
• We choose the examples, the model, how to score a miss, and how big a step to take.

Tensors

Suppose you have three inputs: 1, 2, 3. PyTorch keeps them in a tensor: a block of numbers you can calculate with all at once. The difference from a Python list shows up as soon as you do arithmetic:

A tensor’s arithmetic applies to every number inside it.

import torch

inputs = torch.tensor([[1.0], [2.0], [3.0]])
print(inputs * 2)              # three rows containing 2, 4, 6
print(inputs.shape)            # torch.Size([3, 1])
print(inputs[0])               # tensor([1.]) — the first row
print(inputs[0, 0].item())     # 1.0 — one plain Python number

"""Run with: python3 examples/pytorch_tensors.py"""
import torch

inputs = torch.tensor([[1.0], [2.0], [3.0]])
print([1.0, 2.0, 3.0] * 2)  # a Python list repeats itself
print(inputs * 2)           # a tensor doubles every number
print(inputs.shape)
print(inputs[0])
print(inputs[0, 0].item())
print(inputs.dtype, inputs.device)  # the defaults: float32, on the CPU

Shape counts the entries along each direction. In this lesson each row is one example, and each column is one number describing it:

PyTorch layers expect the second kind: rows of examples.

Every tensor also has a dtype, its number format, and a device, where it lives (CPU or GPU). The defaults, float32 on the CPU, are what this lesson uses; float32 allows fractions such as 1.6.

Linear model

Our model multiplies the input by a weight, then adds a bias:

\[ \text{guess} = \text{weight} \times \text{input} + \text{bias} \]

Drawn as a line, the weight is how steep the line is, and the bias is where it crosses input 0. We start with weight 1 and bias 0:

Learning will tilt and lift the line until it passes through the wanted point.

In PyTorch this model is nn.Linear(1, 1). The two numbers are how many numbers go in and how many come out: one input, one output.

With one of each, the layer holds exactly one weight and one bias. That match is a coincidence: a layer has one weight for every input-output pair, and one bias per output. Bigger layers, in Neural network: one input, have many more.

Module

an object that holds a calculation together with its adjustable numbers. nn.Linear is one.

Parameters

a module’s adjustable numbers. model.parameters() hands them to whatever will adjust them.

Calling model(input_value) returns the guess as a tensor.

Loss

The guess 2 is three short of 5. We turn the miss into one score by squaring it: \((2 - 5)^2 = 9\).

Loss

one number that scores how wrong the model is. Lower is better.

Why square the miss?

• It makes a miss in either direction positive: a guess of 8 also scores 9.
• It makes big misses count far more than small ones: a guess of 4 scores only 1.
• A perfect 5 scores 0.

Plot the score for every possible guess and you get a bowl:

Learning means moving the guess down into the bowl.

PyTorch’s nn.MSELoss computes this score. MSE stands for mean squared error: the average of the squared misses.

With a single example, the average is just the square: 9.

Gradient

From guess 2, moving right goes downhill. But we can’t move the guess directly: it comes from the weight and the bias. So we ask each of them, “if I nudge you up a little, what happens to the loss?”

The loss changes about −0.12 for a nudge of 0.01: an estimated rate of −12.

A smaller nudge gets closer to the slope right at the starting point.

Gradient

the slope of the loss for one parameter: how fast the loss changes as that parameter moves. The weight’s gradient here is −12.

It is negative because raising the weight lowers the loss.

The bias’s gradient is only −6. Nudging the bias by 0.01 moves the guess by 0.01; nudging the weight moves it by 0.02, because the weight is multiplied by the input 2.

• loss.backward() computes these rates for every parameter at once, exactly, without trying any nudges.
• It writes each one beside its parameter, in a slot called .grad.
• It does not change the parameters.

The optimizer makes the move. Plain SGD subtracts \(\text{learning rate} \times \text{gradient}\) from each parameter. (What “stochastic” means is explained with batches in Gradients and backpropagation.) With a learning rate of 0.05:

We choose the learning rate. The gradient only describes the slope where we are standing, so a step that is too big can jump right past the bottom of the bowl. Try it below.

The Gradients and backpropagation lesson slows this step down: why each parameter gets its own direction, the chain rule, and how a batch averages its examples.

Update steps

In code, the update is a short sequence of calls, and the order matters.

Step through it with input 2, answer 5 and learning rate 0.05. Each call changes one thing: watch which box changes.

Code

Here is the whole update as it appears in the example file:

guess_before = model(input_value)
loss_before = loss_fn(guess_before, answer)
optimizer.zero_grad()
loss_before.backward()  # write gradients
optimizer.step()        # change weight and bias

"""A complete PyTorch update with one input and one adjustable model.

Run: python3 examples/pytorch_first_update.py
The matching browser lesson starts from these exact numbers.
"""

import torch
from torch import nn


# The model is guess = weight × input + bias.
model = nn.Linear(1, 1)
with torch.no_grad():
    model.weight.fill_(1.0)
    model.bias.fill_(0.0)

input_value = torch.tensor([[2.0]])
answer = torch.tensor([[5.0]])
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

guess_before = model(input_value)
loss_before = loss_fn(guess_before, answer)
optimizer.zero_grad()  # Remove notes from an earlier example, if any.
loss_before.backward()  # Write a gradient beside each parameter; do not move it yet.

print(f"Guess before: {guess_before.item():.2f}; answer: {answer.item():.2f}")
print(f"Squared error: {loss_before.item():.2f}")
print(f"Gradient for weight: {model.weight.grad.item():.2f}")
print(f"Gradient for bias: {model.bias.grad.item():.2f}")

optimizer.step()  # This is the line that moves the parameters.
with torch.no_grad():
    guess_after = model(input_value)
    loss_after = loss_fn(guess_after, answer)

print(f"Weight: {model.weight.item():.2f}; bias: {model.bias.item():.2f}")
print(f"Guess after: {guess_after.item():.2f}; squared error: {loss_after.item():.2f}")

zero_grad() is there because backward() adds to .grad instead of replacing it. Without clearing, the previous update’s gradients would pile up on top of the new ones.

To check the new guess, the file calls the model inside with torch.no_grad():. That tells PyTorch not to record anything for gradients, because here we only want the number.

Training loop

One example is easy to fit, because many different lines pass through a single point. Real training uses many examples.

Give the model three inputs, 1, 2, 3, with answers 3, 5, 7, stacked as rows: shape (3, 1).

Batch

several examples stacked into one tensor, one example per row.

The same single weight and bias make all three guesses at once. The loss averages the three squared misses, and backward() finds one combined direction. Then we repeat the update, which is all a training loop is:

inputs = torch.tensor([[1.0], [2.0], [3.0]])
answers = torch.tensor([[3.0], [5.0], [7.0]])

for step in range(500):
    guesses = model(inputs)
    loss = loss_fn(guesses, answers)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

"""Repeat the update on three examples until the model fits them.

Run: python3 examples/pytorch_training_loop.py
The browser lesson's training animation uses the same numbers.
"""

import torch
from torch import nn


# Three examples, one per row. The answers follow 2 × input + 1.
inputs = torch.tensor([[1.0], [2.0], [3.0]])
answers = torch.tensor([[3.0], [5.0], [7.0]])

model = nn.Linear(1, 1)
with torch.no_grad():
    model.weight.fill_(1.0)
    model.bias.fill_(0.0)

loss_fn = nn.MSELoss()  # The average of the three squared misses.
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for step in range(500):
    guesses = model(inputs)
    loss = loss_fn(guesses, answers)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if step in (0, 9, 49, 499):
        print(f"Step {step + 1:3d}: loss {loss.item():.4f}  "
              f"weight {model.weight.item():.3f}  bias {model.bias.item():.3f}")

with torch.no_grad():
    print("A new input, 10, gives", round(model(torch.tensor([[10.0]])).item(), 2))

This training run uses a learning rate of 0.01, so the gradual change is easier to follow. Press play.

• Loss falls from 9.67 to about 3.00 after five updates,
• and to 0.12 after twenty.

The line moves toward the three points, then improves more slowly as the weight and bias settle.

After 500 updates the model is \(\text{guess} \approx 2.062 \times \text{input} + 0.859\), with loss about 0.00285.

It is approaching the rule behind the examples, \(2 \times \text{input} + 1\), but has not reached it yet.

For a new input, 10, it predicts about 21.48; the wanted answer is 21. In real projects you keep some examples aside, never used for updates, to test exactly that.

Recap

Here is the training workflow in code order: set up the examples, model and optimizer once, then repeat the five update operations 500 times.

• inputs holds the numbers we give the model: 1, 2, 3, stored as three rows in a tensor.
• answers holds the wanted outputs: 3, 5, 7. Each row matches the same row in inputs. Multiple examples help pin down the model; one pair such as 2 → 5 fits many possible weights and biases.
• model = nn.Linear(1, 1) creates the calculation \(\text{weight} \times \text{input} + \text{bias}\). We start its weight at 1 and bias at 0; training adjusts those same parameters throughout the loop.
• loss_fn = nn.MSELoss() creates the scoring function. It squares each prediction’s gap from its answer, then averages those squared gaps into one loss.
• optimizer = torch.optim.SGD(model.parameters(), lr=0.01) creates the updater. It keeps access to the model’s parameters and uses learning rate 0.01 to scale each adjustment.
• for step in range(500) repeats the following five operations. The model and optimizer were created before the loop, so each iteration carries forward the updated weight and bias.
• guesses = model(inputs) calculates fresh predictions for all three examples using the current weight and bias.
• loss = loss_fn(guesses, answers) compares those predictions with the wanted answers and calculates the current loss.
• optimizer.zero_grad() clears gradients left from the previous update, so this update starts with fresh gradient slots.
• loss.backward() calculates how sensitive this loss is to each parameter and writes those gradients into their .grad slots. The parameters have not moved yet.
• optimizer.step() changes the weight and bias using those gradients and the learning rate. The next iteration calls the model again to see what that adjustment achieved.

After training, model(new_inputs) uses the resulting weight and bias to predict answers. Calling the model makes predictions; the update sequence is what teaches it.