🎓 A1Learner
🔍

PyTorch Pretrain on a Synthetic Task

Pretraining is training a model on a first task so its layers learn reusable features. This lesson pretrains the shapes CNN on shape naming and saves the extractor.

Three shapes as pixels

Task A data is the shapes data. Each image is 16 by 16 pixels with values in 0 to 1.

x, y = make_shapes(30, seed=1)
for cls, name in ((0, "bar"), (1, "ring"), (2, "square")):
    first = x[y == cls][0]
    print(name, "filled pixels:", int(first.sum()), "shape:", tuple(first.shape))

Training is the usual loop

Pretraining uses the same loop as any training run. Only the plan for the model changes.

for epoch in range(30):
    model.train()
    opt.zero_grad()
    loss = loss_fn(model(train_x), train_y)
    loss.backward()
    opt.step()

Saving only the extractor

A state_dict holds a layer's learned tensors. Saving the extractor alone keeps the classifier out of the file.

torch.save(model.features.state_dict(), "features.pt")
print("Extractor tensors:", list(model.features.state_dict().keys()))

Example

import torch
from torch import nn

torch.manual_seed(0)


def make_shapes(per_class, seed):
    g = torch.Generator().manual_seed(seed)
    ys, xs = torch.meshgrid(torch.arange(16), torch.arange(16), indexing="ij")
    images, labels = [], []
    for cls in range(3):
        for _ in range(per_class):
            img = torch.zeros(16, 16)
            if cls == 0:  # bar: a filled horizontal band
                top = int(torch.randint(2, 12, (1,), generator=g))
                img[top:top + 3, 2:14] = 1.0
            elif cls == 1:  # ring: a band at one distance from the centre
                cy = 8 + int(torch.randint(-1, 2, (1,), generator=g))
                cx = 8 + int(torch.randint(-1, 2, (1,), generator=g))
                dist = torch.sqrt((ys - cy).float() ** 2 + (xs - cx).float() ** 2)
                img[(dist >= 4.5) & (dist <= 6.5)] = 1.0
            else:  # square: a filled 8 by 8 block
                top = int(torch.randint(0, 9, (1,), generator=g))
                left = int(torch.randint(0, 9, (1,), generator=g))
                img[top:top + 8, left:left + 8] = 1.0
            images.append(img)
            labels.append(cls)
    x = torch.stack(images).unsqueeze(1)
    y = torch.tensor(labels)
    perm = torch.randperm(len(y), generator=g)
    return x[perm], y[perm]


class ShapesCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 8, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(8, 16, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Linear(16, 3)

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        return self.classifier(x)


train_x, train_y = make_shapes(100, seed=1)
test_x, test_y = make_shapes(50, seed=2)
print("Task A train rows:", train_x.shape[0])

model = ShapesCNN()
opt = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(30):
    model.train()
    opt.zero_grad()
    loss = loss_fn(model(train_x), train_y)
    loss.backward()
    opt.step()

model.eval()
with torch.no_grad():
    pred = model(test_x).argmax(dim=1)
acc = float((pred == test_y).float().mean())
print(f"Task A test accuracy: {acc:.4f}")

torch.save(model.features.state_dict(), "features.pt")
saved = torch.load("features.pt")
print("Saved tensors:", len(saved))
print("Extractor saved to features.pt")

Key points

  • Task A: the first task, trained properly so the features are worth reusing
  • Use torch.save with a state_dict to store the extractor's learned tensors
  • The saved file holds 4 tensors, the weights and biases of two conv layers
  • A fixed seed makes the shapes data identical on every run

In this example, make_shapes builds task A with a seeded generator. The training set holds 300 rows, and the test set holds 150 rows. A fresh shapes CNN trains on task A for 30 epochs with Adam. The test accuracy prints as 1.0000, so the features are learned properly. Then torch.save writes the extractor's state_dict to features.pt. The program loads the file back and counts 4 saved tensors. The next lesson loads these tensors into a fresh model.

Note

The extractor is saved without its classifier, because the head belongs to task A alone.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why save only the feature extractor after pretraining?