Convolutional networks

Overview

The task: a ball on a 16 × 16 screen, and a network that says where its centre is, in pixels. Training images never show the ball in the bottom-right corner, but the test does.

The mental model: a convolution looks at every position with the same small set of weights, so what it learns in one place works in every place.

Image

To PyTorch an image is a tensor of brightness values with three sizes:

Channels

how many numbers each pixel has: 1 for grey, 3 for a colour camera (red, green, blue).

Height and width

the rows and columns of pixels.

Our image is (1, 16, 16): one channel, 16 rows, 16 columns. A batch adds a fourth size in front, so 64 images are (64, 1, 16, 16). One of ACT’s camera frames is (3, 480, 640).

def draw(centres):
    """One image per (x, y) centre: brightness 1 inside the ball, 0 outside, a soft edge one pixel wide."""
    distance = ((COLS - centres[:, 0, None, None]) ** 2 + (ROWS - centres[:, 1, None, None]) ** 2).sqrt()
    return (RADIUS + 0.5 - distance).clamp(0, 1).unsqueeze(1)          # shape (N, 1, 16, 16): N images, 1 channel

"""How a network reads an image: a ball on a 16 × 16 screen, and three networks that report where it is.

Run: python3 site/ml/foundations/examples/pytorch_cnn.py

Training images never show the ball in the bottom-right corner (centre x > 9 and y > 9). The test asks for
positions there too. All three networks train with Adam.
"""

import torch
from torch import nn

SIZE, RADIUS = 16, 2.0
ROWS, COLS = torch.meshgrid(torch.arange(SIZE) + 0.5, torch.arange(SIZE) + 0.5, indexing="ij")   # pixel centres


def draw(centres):
    """One image per (x, y) centre: brightness 1 inside the ball, 0 outside, a soft edge one pixel wide."""
    distance = ((COLS - centres[:, 0, None, None]) ** 2 + (ROWS - centres[:, 1, None, None]) ** 2).sqrt()
    return (RADIUS + 0.5 - distance).clamp(0, 1).unsqueeze(1)          # shape (N, 1, 16, 16): N images, 1 channel


def held_out(centres):
    return (centres[:, 0] > 9) & (centres[:, 1] > 9)                   # the bottom-right corner


def images(count, seed, corner=False):
    """`count` images with the ball anywhere except the corner, or (corner=True) only in the corner."""
    generator, kept = torch.Generator().manual_seed(seed), []
    while sum(len(k) for k in kept) < count:
        centres = 2 + 12 * torch.rand(4 * count, 2, generator=generator)
        kept.append(centres[held_out(centres) == corner])
    centres = torch.cat(kept)[:count]
    return draw(centres), centres


class SpatialSoftmax(nn.Module):
    """Softmax over the 256 pixels of a one-channel map, then the average pixel position under those weights."""
    def forward(self, maps):
        weights = maps.flatten(1).softmax(dim=1).view(-1, SIZE, SIZE)
        return torch.stack([(weights * COLS).sum((1, 2)), (weights * ROWS).sum((1, 2))], dim=1)


def make_network(kind):
    if kind == "dense":            # every pixel has its own weights
        return nn.Sequential(nn.Flatten(), nn.Linear(SIZE * SIZE, 3), nn.ReLU(), nn.Linear(3, 2))
    if kind == "cnn":              # shared 3 × 3 filters, pooling, then a dense layer for the answer
        return nn.Sequential(nn.Conv2d(1, 4, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
                             nn.Conv2d(4, 4, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
                             nn.Flatten(), nn.Linear(4 * 4 * 4, 8), nn.ReLU(), nn.Linear(8, 2))
    return nn.Sequential(          # "heatmap": shared filters only, and a fixed formula reads the map
        nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.Conv2d(8, 8, 3, padding=1), nn.ReLU(),
        nn.Conv2d(8, 1, 3, padding=1), SpatialSoftmax())


LEARNING_RATE = {"dense": 0.01, "cnn": 0.01, "heatmap": 0.03}     # Adam rates chosen on validation images


def train(kind, seed=0, lr=None, epochs=60):
    lr = lr or LEARNING_RATE[kind]
    torch.set_num_threads(1)
    torch.manual_seed(seed)
    pictures, centres = images(2000, seed)
    network = make_network(kind)
    optimizer = torch.optim.Adam(network.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    order = torch.Generator().manual_seed(seed)
    for epoch in range(epochs):
        for batch in torch.randperm(len(pictures), generator=order).split(64):
            loss = loss_fn(network(pictures[batch]), centres[batch])     # squared miss in pixels
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    return network


def miss(network, data):
    """Average distance, in pixels, between the predicted and the true centre."""
    pictures, centres = data
    with torch.no_grad():
        return (network(pictures) - centres).norm(dim=1).mean().item()


if __name__ == "__main__":
    pictures, centres = images(1, seed=5)
    print("one image: tensor shape", tuple(pictures.shape), "ball centre", [round(v, 2) for v in centres[0].tolist()])

    # One learned 3 × 3 filter over a 5 × 5 patch around the ball, by hand and by nn.Conv2d.
    heatmap = train("heatmap")
    conv = nn.Conv2d(1, 1, 3)                                         # no padding: a 5 × 5 patch gives 3 × 3
    with torch.no_grad():
        conv.weight.copy_(heatmap[0].weight[:1])                       # the first layer's filter 0
        conv.bias.copy_(heatmap[0].bias[:1])
        patch = pictures[0, 0, 1:6, 9:14]                              # rows 1-5, columns 9-13
        by_torch = conv(patch[None, None])[0, 0]
    by_hand = [[(patch[r:r + 3, c:c + 3] * conv.weight[0, 0]).sum().item() + conv.bias.item() for c in range(3)] for r in range(3)]
    print("patch:", [[round(v, 3) for v in row] for row in patch.tolist()])
    print("filter:", [[round(v, 3) for v in row] for row in conv.weight[0, 0].tolist()], "bias", round(conv.bias.item(), 3))
    print("by hand: ", [[round(v, 4) for v in row] for row in by_hand])
    print("by torch:", [[round(v, 4) for v in row] for row in by_torch.tolist()])

    for kind in ("dense", "cnn", "heatmap"):
        print(f"{kind:8s}: {sum(p.numel() for p in make_network(kind).parameters()):5d} parameters")
    print("a dense layer from 256 pixels to 256 outputs:", sum(p.numel() for p in nn.Linear(256, 256).parameters()),
          "parameters; a 3 × 3 convolution from 1 to 8 channels:", sum(p.numel() for p in nn.Conv2d(1, 8, 3).parameters()))

    test, corner = images(500, seed=77), images(500, seed=78, corner=True)
    for kind in ("dense", "cnn", "heatmap"):
        network = train(kind)
        print(f"{kind:8s} seed 0: miss {miss(network, test):.3f} px on usual positions, {miss(network, corner):.3f} px in the unseen corner")

• distance is each pixel’s distance from the ball’s centre.
• clamp(0, 1) makes pixels within the 2-pixel radius bright, with a one-pixel soft edge, so a centre between pixels still shows.
• unsqueeze(1) adds the channel size.

Training uses 2,000 images; the ball’s centre is anywhere from 2 to 14 pixels in each direction, except the corner where both are above 9.

Dense on pixels

The first network reads the image the way Catch did: flatten the 256 pixels into one list, then a dense layer. Here 256 → 3 ReLU → 2, which is 779 weights and biases.

A dense layer gives every pixel its own weight into every neuron. During training, a weight’s gradient is its pixel’s value times the error signal reaching that neuron. A pixel that is 0 in every training image therefore gets a gradient of exactly 0, so its weights never move.

• Blue is a positive weight and magenta a negative one; stronger colour is a larger weight.
• The outlined pixels were 0 in every training image.

Those pixels still hold exactly their random starting weights: the trained network has learned nothing about them. A ball in the corner lights exactly these pixels.

Convolution

A convolution layer instead has a few small filters, here 3 × 3 weights each. Each filter slides over the image. At every position it multiplies its 9 weights by the 9 pixels under it, adds them up and adds a bias.

Here is one learned filter on a 5 × 5 patch around a ball. The 3 × 3 filter fits in 3 × 3 positions, so the output is 3 × 3:

The same numbers come from nn.Conv2d with this filter’s weights, to every printed digit. The example checks that.

Convolution layer

small filters slid over the whole image; each position gives the sum of filter weights times pixels, plus a bias.

Filter (kernel)

one small grid of weights, the pattern a convolution looks for.

return nn.Sequential(          # "heatmap": shared filters only, and a fixed formula reads the map
    nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.Conv2d(8, 8, 3, padding=1), nn.ReLU(),
    nn.Conv2d(8, 1, 3, padding=1), SpatialSoftmax())

• nn.Conv2d(1, 8, 3, padding=1): 1 channel in, 8 filters out, 3 × 3 each. padding=1 adds a border of zeros, so the output stays 16 × 16.
• Three convolutions end in one 16 × 16 map: high where the network thinks the ball is.
• SpatialSoftmax turns that map into a position. Softmax over the 256 pixels gives weights that add up to 1, as softmax over the three sizes did in the several-inputs lesson. The position is the weighted average of the pixel positions.

Feature maps

Each filter’s output over the whole image is a feature map: the image re-drawn as “how much of my pattern is here”. The first layer’s eight filters, and their eight maps for one image:

• Filters: blue weights positive, magenta negative.
• Maps: brighter is a larger output after ReLU; black is zero.

Each filter has learned a different pattern. Some maps light up on the ball, one on a single edge of it, and some on the empty background, with a dark hole where the ball is. The next layer combines the eight maps.

Feature map

the output of one filter over the whole image: where its pattern appears, and how strongly.

Weight sharing

The filter at the top-left of the image and the one at the bottom-right are the same nine weights. That has two consequences.

• Far fewer parameters. A 3 × 3 convolution from 1 to 8 channels has weights and biases. A dense layer from 256 pixels to 256 outputs has .
• The same answer everywhere. A filter that has learned what a ball’s edge looks like finds it in any position, including ones training never showed.

Weight sharing

using the same weights at every position of the input.

Pooling

Most image networks also shrink the maps as they go, so that later layers see a larger area. Our second network does that:

return nn.Sequential(nn.Conv2d(1, 4, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
                     nn.Conv2d(4, 4, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
                     nn.Flatten(), nn.Linear(4 * 4 * 4, 8), nn.ReLU(), nn.Linear(8, 2))

• nn.MaxPool2d(2) keeps the largest of every 2 × 2 block, so 16 × 16 becomes 8 × 8, then 4 × 4.
• Stride does the same inside a convolution: with stride 2 the filter jumps two pixels at a time. ResNet-18 uses both.
• The last part is dense again: 4 channels × 4 × 4 cells = 64 numbers into 8 neurons, then 2 outputs.

Pooling

shrinking a feature map by keeping one number per block, usually the largest.

Stride

how many pixels a filter moves between positions.

Test

All three networks have about the same budget: weights and biases. Each trains with Adam on the same 2,000 images, at the rate chosen for it on 500 validation images.

Predict before you click: which network finds a ball in the corner it never saw?

ACT’s backbone

ACT reads its camera frames with ResNet-18, a convolutional network of the same kind, only far bigger: 18 layers deep, with weights in the part ACT uses. In LeRobot 0.6.1:

• Backbone: resnet18 from torchvision, without its final classifier.
• Pretrained: it starts from ResNet18_Weights.IMAGENET1K_V1, weights already trained on ImageNet’s 1.28 million labelled photos.
• Output: a frame becomes a feature map. Stride and pooling shrink it 32 times in each direction.
• Tokens: each of the 15 × 20 = 300 cells, 512 numbers, becomes one input to ACT’s transformer, the next lesson’s topic.
• Training: the backbone keeps learning with the rest of ACT, at the same rate \(10^{-5}\).

Why start from ImageNet? Filters for edges, corners and textures take millions of images to learn, and a robot’s 50 demonstrations hold far too few. The pretrained filters already find such patterns anywhere in a frame, so ACT’s own training only has to adjust them.

Backbone

the part of a network that turns raw input, here an image, into feature maps for the rest of the model.

Pretrained

starting from weights already trained on another, larger task.

Recap

• draw(centres) makes images as tensors: (N, 1, 16, 16).
• make_network("dense") gives each pixel its own weights; pixels that never light up in training never learn.
• nn.Conv2d(1, 8, 3, padding=1) slides 8 shared 3 × 3 filters over the image and makes 8 feature maps.
• nn.MaxPool2d(2) halves a map’s height and width.
• SpatialSoftmax() reads the final map as a position: softmax over pixels, then the average position.
• miss(network, data) grades each network on usual positions and in the unseen corner.

A convolution is a dense layer that looks at a small patch and reuses its weights everywhere. That is why it needs few parameters, and why what it learns in one place works in another.