PyTorch Freeze and Reuse
Freezing a layer sets requires_grad to False so training leaves its weights unchanged. A frozen extractor keeps task A's features fixed while a new head learns task B.
Freezing the extractor
One loop freezes every extractor parameter. The head keeps requires_grad set to True.
for p in model.features.parameters():
p.requires_grad = False
print("Extractor trainable:", [p.requires_grad for p in model.features.parameters()])
print("Head trainable:", [p.requires_grad for p in model.classifier.parameters()])Task B is the same shapes, rotated
torch.rot90 turns each image by 90 degrees. The bar that lay flat now stands tall, and the pixel count stays the same.
img = torch.zeros(1, 1, 16, 16)
img[0, 0, 4:7, 2:14] = 1.0
rot = torch.rot90(img, 1, dims=(2, 3))
print("Filled before:", int(img.sum()), "after:", int(rot.sum()))Only the head reaches the optimizer
A frozen extractor is a fixed feature maker. Only the head adapts to the new task.
trainable = [p for p in frozen.parameters() if p.requires_grad]
opt = torch.optim.Adam(trainable, lr=0.01)
print("Optimizer parameters:", sum(p.numel() for p in trainable))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)
def train_adam(model, x, y, epochs, lr=0.01):
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 the full CNN 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)
# Task B: the same shapes rotated 90 degrees. Only 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])
# The frozen model: the extractor stays exactly as task A left it.
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
trainable = sum(p.numel() for p in frozen.parameters() if p.requires_grad)
print("Frozen model trainable parameters:", trainable)
train_adam(frozen, rows_b, labels_b, epochs=20)
print(f"Frozen extractor accuracy: {accuracy(frozen, test_b, labels_test_b):.4f}")
# The scratch model: a fresh CNN learns from the same 30 rows alone.
torch.manual_seed(0)
scratch = ShapesCNN()
total = sum(p.numel() for p in scratch.parameters())
print("Scratch model trainable parameters:", total)
train_adam(scratch, rows_b, labels_b, epochs=20)
print(f"From-scratch accuracy: {accuracy(scratch, test_b, labels_test_b):.4f}")Key points
- Freezing: setting requires_grad to False so the optimizer leaves a layer's weights alone
- Use a fresh head for the new task, and train only that head
- Trainable parameters in the frozen model: 51, the head alone
- Trainable parameters in a fresh CNN: 1,299
- Both runs in the closing example train for 20 epochs on the same 30 rows
In this example, the program first pretrains a shapes CNN on task A's 300 rows, as the pretraining lesson did. Task B is the same shapes rotated 90 degrees. The program takes 10 labelled rows per class, so task B has 30 rows in all. That is one tenth of task A's rows. The frozen model loads the pretrained extractor and freezes it, so only its 51 head parameters train. It reaches 0.9467 on the task B test set after 20 epochs. A fresh CNN trains on the same 30 rows for the same 20 epochs. It reaches 0.6667. The reused features carry task A's training into task B, and the small head needs few rows to read them. The next lesson loosens the freeze and fine-tunes the last block.
Note
requires_grad is a flag on each parameter tensor, and freezing changes training, not the forward pass.
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
A model has a frozen extractor and a fresh head. Training runs for 20 epochs. What changes?