PyTorch CNNs and Transfer Learning for Image Classification

CNN Shape Flow

Convolutional neural networks learn spatial patterns in images. PyTorch and torchvision make it practical to build CNNs from scratch or adapt pretrained models. In most production image tasks, transfer learning is the fastest strong baseline.

Transfer learning uses a model pretrained on a large dataset, replaces the final classifier, and fine-tunes it for your classes. This saves data, compute, and development time.

Images usually enter as `[batch, channels, height, width]`. Convolution layers preserve spatial structure, pooling reduces spatial dimensions, and classifier layers map features to class scores.

Use transforms to resize, normalize, and augment images.
Match normalization values expected by pretrained weights.
Freeze early layers for small datasets, then optionally fine-tune later layers.

Fine-Tune a Pretrained ResNet

This example replaces the classification head for a custom number of classes.

Fine-Tune a Pretrained ResNet

import torch
from torch import nn
from torchvision.models import resnet18, ResNet18_Weights

num_classes = 5
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

weights = ResNet18_Weights.DEFAULT
model = resnet18(weights=weights)

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

in_features = model.fc.in_features
model.fc = nn.Sequential(
    nn.Dropout(p=0.2),
    nn.Linear(in_features, num_classes),
)

model = model.to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.fc.parameters(), lr=1e-3, weight_decay=1e-4)

print(weights.transforms())
print(model.fc)

Train the new head first. Then unfreeze some deeper layers for fine-tuning if needed.
Use the transforms associated with pretrained weights for correct preprocessing.

Freeze a Feature Extractor Before Creating the Optimizer

import torch

features = torch.nn.Sequential(torch.nn.Linear(8, 6), torch.nn.ReLU())
classifier = torch.nn.Linear(6, 3)
for parameter in features.parameters():
    parameter.requires_grad = False

model = torch.nn.Sequential(features, classifier)
trainable = [parameter for parameter in model.parameters() if parameter.requires_grad]
optimizer = torch.optim.Adam(trainable, lr=1e-3)
print(len(trainable), len(optimizer.param_groups[0]['params']))

Output

2 2

Before you move on