python 15 lines · 1 tab

Transfer learning with pretrained torchvision backbones

1 tab
import torch.nn as nn
from torchvision.models import resnet50, ResNet50_Weights

model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)

for parameter in model.parameters():
    parameter.requires_grad = False

model.fc = nn.Sequential(
    nn.Dropout(0.3),
    nn.Linear(model.fc.in_features, 4),
)

for parameter in model.fc.parameters():
    parameter.requires_grad = True
1 file · python Explain with highlit

Transfer learning is the right default when labeled data is limited and time matters. I usually freeze the backbone first, train the head, then selectively unfreeze deeper layers if the domain gap justifies it. This strategy converges faster and is much less brittle than training from scratch.

Share this code

Here's the card — post it anywhere.

Transfer learning with pretrained torchvision backbones — share card
Link copied