PyTorch Fine-Tuning

The extractor freezes first. Then the last conv block, at index 3, unfreezes again.

for p in model.features.parameters():
    p.requires_grad = False
for p in model.features[3].parameters():
    p.requires_grad = True
print("Conv 1 trainable:", [p.requires_grad for p in model.features[0].parameters()])
print("Conv 2 trainable:", [p.requires_grad for p in model.features[3].parameters()])
print("Head trainable:", [p.requires_grad for p in model.classifier.parameters()])

The rate here is 0.003, not the 0.01 used for pretraining. A small rate nudges the features instead of replacing them.

trainable = [p for p in model.parameters() if p.requires_grad]
opt = torch.optim.Adam(trainable, lr=0.003)
print("Optimizer parameters:", sum(p.numel() for p in trainable))

Both runs start from the same pretrained extractor and the same 30 rows. Only the unfrozen block and the rate differ.

print(f"Frozen accuracy:     {frozen_acc:.4f}")
print(f"Fine-tuned accuracy: {tuned_acc:.4f}")

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)


def train_adam(model, x, y, epochs, lr):
    opt = torch.optim.Adam(
        [p for p in model.parameters() if p.requires_grad], lr=lr
    )
    loss_fn = nn.CrossEntropyLoss()
    for _ in range(epochs):
        model.train()
        opt.zero_grad()
        loss = loss_fn(model(x), y)
        loss.backward()
        opt.step()
    return model


def accuracy(model, x, y):
    model.eval()
    with torch.no_grad():
        pred = model(x).argmax(dim=1)
    return float((pred == y).float().mean())


# Task A: pretrain on 300 rows, as in the pretraining lesson.
train_a, labels_a = make_shapes(100, seed=1)
torch.manual_seed(0)
pretrained = ShapesCNN()
train_adam(pretrained, train_a, labels_a, epochs=30, lr=0.01)

# Task B: rotated shapes, 30 labelled rows.
all_b, labels_b_all = make_shapes(100, seed=11)
all_b = torch.rot90(all_b, 1, dims=(2, 3))
test_b, labels_test_b = make_shapes(50, seed=12)
test_b = torch.rot90(test_b, 1, dims=(2, 3))

rows_b, labels_b = [], []
for cls in range(3):
    take = (labels_b_all == cls).nonzero()[:10].flatten()
    rows_b.append(all_b[take])
    labels_b.append(labels_b_all[take])
rows_b = torch.cat(rows_b)
labels_b = torch.cat(labels_b)
print("Task B train rows:", rows_b.shape[0])

# Run 1: the extractor stays frozen. Only the head trains.
torch.manual_seed(0)
frozen = ShapesCNN()
frozen.features.load_state_dict(pretrained.features.state_dict())
for p in frozen.features.parameters():
    p.requires_grad = False
train_adam(frozen, rows_b, labels_b, epochs=60, lr=0.01)
frozen_acc = accuracy(frozen, test_b, labels_test_b)

# Run 2: fine-tuning. The last conv block unfreezes and trains at a small rate.
torch.manual_seed(0)
tuned = ShapesCNN()
tuned.features.load_state_dict(pretrained.features.state_dict())
for p in tuned.features.parameters():
    p.requires_grad = False
for p in tuned.features[3].parameters():
    p.requires_grad = True
trainable = sum(p.numel() for p in tuned.parameters() if p.requires_grad)
print("Fine-tuned model trainable parameters:", trainable)
train_adam(tuned, rows_b, labels_b, epochs=60, lr=0.003)
tuned_acc = accuracy(tuned, test_b, labels_test_b)

print(f"Frozen accuracy:     {frozen_acc:.4f}")
print(f"Fine-tuned accuracy: {tuned_acc:.4f}")