Image Classification in Practice: Transfer Learning and Augmentation

A pretrained image classifier provides a useful starting point when labeled data or training compute is limited. Its backbone converts pixels into features; its head maps those features to class scores. We will replace the ImageNet head with a 37-class head, control which parts learn, and build the training and evaluation steps around that choice. Whether the pretrained features help on a new domain is something to measure.

Choosing what to train

There is no image-count threshold that determines the best regime. Label quality, task similarity, augmentation, regularization, and training duration all affect the comparison. A head-only run is a relatively inexpensive baseline; compare it with partial or full fine-tuning on the same validation split.

For a domain unlike ImageNet, check the input channels and whether the pretrained features distinguish the classes you need. Greater domain difference can motivate broader adaptation or a different pretrained backbone, but it does not determine how many layers to unfreeze. For grayscale inputs, repeating a channel provides a three-channel input; it does not establish that the representation is suitable.

Distribution mismatch and transfer learning are treated separately in Data Mismatch, Transfer Learning, and Multitask Learning.

Replacing the head

import torch
import torch.nn as nn
import torchvision

WEIGHTS = torchvision.models.ResNet50_Weights.IMAGENET1K_V2

def build_model(num_classes, freeze_backbone=True, weights=WEIGHTS):
    model = torchvision.models.resnet50(weights=weights)
    if freeze_backbone:
        for p in model.parameters():
            p.requires_grad = False
    model.fc = nn.Linear(model.fc.in_features, num_classes)
    return model

torch.manual_seed(0)
model = build_model(37)
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(trainable, total)
# 75813 23583845

Run the Python blocks in order with compatible PyTorch and Torchvision installations. The first model construction downloads the selected weights if they are not cached. ResNet-50 supplies 2,048 features per image, so the new head has \(2048\times37+37=75{,}813\) parameters, including biases. nn.Linear creates trainable parameters by default. Freezing before replacing fc therefore leaves the head trainable; alternatively, explicitly freeze only the backbone after replacement.

For this ResNet, freezing parameter gradients does not freeze BatchNorm running means and variances. Those buffers update during training-mode forward passes. The helper below keeps the stored statistics fixed; it does not change whether the affine scale and offset are trainable. With a frozen backbone those parameters are already frozen. Call the helper after each model.train(), which otherwise puts BatchNorm back in training mode. During full fine-tuning, retaining or adapting the statistics is a separate choice to validate.

def freeze_batchnorm_stats(module):
    for m in module.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.eval()

Layer-wise learning rate decay

A smaller learning rate for earlier stages expresses a preference to preserve more of their pretrained features. It is a candidate to compare with a single learning rate, not a guaranteed improvement. Here the head receives \(10^{-3}\), and each earlier stage receives 0.75 times the next stage’s rate. These values are starting points for validation. Learning-rate decay across layers is distinct from weight decay, which regularizes parameter values.

def layerwise_groups(model, base_lr=1e-3, decay=0.75):
    stages = [(model.conv1, model.bn1), (model.layer1,), (model.layer2,),
              (model.layer3,), (model.layer4,), (model.fc,)]
    groups = []
    for depth, stage in enumerate(reversed(stages)):
        params = [p for module in stage for p in module.parameters()
                  if p.requires_grad]
        if params:
            groups.append({"params": params, "lr": base_lr * decay ** depth})
    assigned = [id(p) for g in groups for p in g["params"]]
    expected = {id(p) for p in model.parameters() if p.requires_grad}
    assert len(assigned) == len(set(assigned))
    assert set(assigned) == expected
    return groups

# Random weights suffice to inspect optimizer membership and rates.
full_model = build_model(37, freeze_backbone=False, weights=None)
for g in layerwise_groups(full_model):
    print(len(g["params"]), round(g["lr"], 6))
# 2 0.001
# 30 0.00075
# 57 0.000563
# 39 0.000422
# 30 0.000316
# 3 0.000237

del full_model

The printed counts are parameter tensors, not individual scalars. The stem group includes both conv1 and bn1; omitting the latter would leave trainable parameters outside the optimizer. The assertions check coverage and uniqueness. If you unfreeze layers after creating a head-only optimizer, rebuild its groups or add the new parameters explicitly. Rebuilding also resets optimizer state unless you deliberately restore it.

Augmentation as a label-preserving contract

An ordinary label-preserving augmentation assumes that the transformed image still supports its original label. A horizontal flip may preserve a pet-species label but change a left-facing versus right-facing label. A vertical flip of an overhead scene depends on whether orientation matters to the task. Color changes and crops can remove the evidence needed to identify a class. Inspect transformed examples from each class before selecting their strength. Mixup and CutMix, below, also change the training targets.

Use a defined evaluation pipeline so successive validation scores are comparable. Here training uses random crops and flips, while validation uses the deterministic preprocessing packaged with the selected weights: resize the shorter edge to 232, center-crop to 224, and normalize. Randomized evaluation is possible when it is part of the intended protocol, but introduces sampling variation; it is not necessarily pessimistic. A center crop can also remove relevant content, so check that it suits the images being evaluated.

from torchvision.transforms import v2
from PIL import Image

train_tf = v2.Compose([
    v2.RandomResizedCrop(224, scale=(0.65, 1.0), antialias=True),
    v2.RandomHorizontalFlip(),
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
val_tf = WEIGHTS.transforms()

# A synthetic RGB image checks the preprocessing interface.
image = Image.new("RGB", (320, 280), color=(80, 120, 160))
print(tuple(train_tf(image).shape), tuple(val_tf(image).shape))
print(torch.equal(val_tf(image), val_tf(image)))
# (3, 224, 224) (3, 224, 224)
# True

Both pipelines expect RGB images. For a raw PIL image, convert it with image.convert("RGB"); Torchvision’s default ImageFolder loader does this for ordinary image files. Matching the pretrained normalization provides a baseline. Different preprocessing can be evaluated deliberately, but accidental differences between training and inference make results hard to interpret. After testing the crop-and-flip baseline, you can compare RandAugment or RandomErasing. RandomErasing with its default fill of zero after normalization fills a region with the channel means in the original pixel scale. The useful strengths depend on the task.

Training the head and evaluating a checkpoint

A data loader supplies a batch shaped (B, 3, H, W) and B integer class indices from 0 to 36. Use the same class-to-index mapping in every split; if you build separate ImageFolder datasets, check their class_to_idx dictionaries. Give the training dataset train_tf and the validation and test datasets val_tf. Split source images before generating augmented views, according to the grouping policy discussed below.

def train_epoch(model, loader, optimizer, device):
    model.train()
    freeze_batchnorm_stats(model)
    criterion = nn.CrossEntropyLoss()
    total_loss, count = 0.0, 0
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad(set_to_none=True)
        logits = model(images)
        loss = criterion(logits, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * labels.numel()
        count += labels.numel()
    if count == 0:
        raise ValueError("empty training loader")
    return total_loss / count

@torch.no_grad()
def evaluate(model, loader, device):
    model.eval()
    total_loss, correct, count = 0.0, 0, 0
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        logits = model(images)
        total_loss += nn.functional.cross_entropy(
            logits, labels, reduction="sum").item()
        correct += (logits.argmax(dim=1) == labels).sum().item()
        count += labels.numel()
    if count == 0:
        raise ValueError("empty evaluation loader")
    return {"loss": total_loss / count, "accuracy": correct / count}

# One small synthetic batch tests updates, not classification quality.
device = torch.device("cpu")
model = model.to(device)
optimizer = torch.optim.SGD(layerwise_groups(model), momentum=0.9)
batch = [(torch.randn(2, 3, 64, 64), torch.tensor([0, 1]))]
head_before = model.fc.weight.detach().clone()
mean_before = model.bn1.running_mean.clone()
train_epoch(model, batch, optimizer, device)
metrics = evaluate(model, batch, device)
print("head changed:", not torch.equal(head_before, model.fc.weight))
print("BatchNorm mean fixed:", torch.equal(mean_before, model.bn1.running_mean))
print("finite evaluation loss:", bool(torch.isfinite(torch.tensor(metrics["loss"]))))
# head changed: True
# BatchNorm mean fixed: True
# finite evaluation loss: True

The example uses the pretrained backbone, but its random 64-by-64 inputs and arbitrary labels only test the mechanics. Real batches from the transforms above are 224-by-224. The model returns 37 logits per image; cross-entropy consumes them directly, without a preceding softmax. Evaluation sums losses and correct predictions across images, so a smaller last batch receives its proper weight. In this head-only run, the backbone parameters and BatchNorm statistics stay fixed while the head changes.

For a real run, call train_epoch with your training loader each epoch and evaluate with the validation loader. Choose a metric in advance and save the checkpoint that improves it, together with the class mapping and preprocessing configuration. Reload that checkpoint for final testing; do not use test results to choose the epoch or augmentation. To use a GPU, move the model to that device before constructing the optimizer; both functions move each batch to it. The example keeps BatchNorm statistics fixed even if layers are later unfrozen, so change that policy deliberately if you want to adapt them.

Mixup and CutMix

Mixup uses \(\tilde x=\lambda x_i+(1-\lambda)x_j\) and mixes the targets by the same fraction. With a cat and a dog and \(\lambda=0.75\), the target is 75% cat and 25% dog. The Beta distribution draws this fraction between 0 and 1; a small positive alpha favors values near the endpoints. CutMix pastes a rectangular patch instead, using the actual pasted area after clipping to the image boundaries to determine the target fraction. Area is a training heuristic, not a measurement of how much class evidence the patch contains.

import numpy as np

def mixup(x, y, alpha=0.2):
    if alpha <= 0:
        raise ValueError("alpha must be positive")
    lam = np.random.beta(alpha, alpha)
    idx = torch.randperm(x.size(0), device=x.device)
    return lam * x + (1 - lam) * x[idx], y, y[idx], lam

def mixup_loss(criterion, logits, y_a, y_b, lam):
    return lam * criterion(logits, y_a) + (1 - lam) * criterion(logits, y_b)

# For unweighted mean cross-entropy, mixing losses equals mixing targets.
logits = torch.tensor([[2., 0.], [0., 2.]])
y_a = torch.tensor([0, 1])
y_b = torch.tensor([1, 0])
lam = 0.75
soft_targets = (lam * torch.nn.functional.one_hot(y_a, 2)
                + (1 - lam) * torch.nn.functional.one_hot(y_b, 2)).float()
loss = mixup_loss(nn.CrossEntropyLoss(), logits, y_a, y_b, lam)
print(round(loss.item(), 6))
print(torch.allclose(loss, nn.CrossEntropyLoss()(logits, soft_targets)))
# 0.626928
# True

The helper uses one mixing fraction per batch and expects images and integer labels on the same device. In the training step, unpack mixed_x, y_a, y_b, lam = mixup(images, labels), run the model on it, and replace the ordinary loss with mixup_loss(criterion, logits, y_a, y_b, lam). Seed both NumPy and PyTorch when reproducing this implementation. The equality above concerns unweighted mean cross-entropy; class weighting can change its reduction denominators. Mixup and CutMix can improve generalization, but gains and training time depend on the data and recipe. They can be combined with label smoothing; compare their joint strength on validation data instead of assuming the combination will underfit.

Class imbalance

A weighted sampler changes which examples appear; class weights in the loss change their contributions; a decision-rule adjustment changes the predictions without retraining. These address different aspects of imbalance. For a binary classifier, a validation-tuned probability threshold can trade precision against recall. A 37-class classifier using argmax has no single binary threshold; class-specific decisions require a defined objective or cost model. Reweighting training may also change the interpretation of output probabilities under deployment class frequencies. Evaluate on a representative validation distribution.

Report per-class recall alongside overall accuracy, including how many examples support each class’s estimate. Predicting only the majority class gives 95% accuracy when that class makes up 95% of the data, while identifying none of the other classes. The confusion matrix shows which classes are being confused; the headline accuracy alone cannot show that pattern.

Metric learning and the embedding-plus-threshold pipeline are covered in Face Recognition and Metric Learning: Siamese Networks, Triplet Loss, ArcFace.

Test-time augmentation and evaluation

Test-time augmentation averages predictions over several allowed views of an image. For example, average the class-probability vectors for the original and a horizontally flipped image, then take the largest component. This uses two forward passes and can help or hurt accuracy. Compare the gain and measured latency on validation data before choosing the final protocol, and apply it consistently to the test set. A flip is a candidate only when it preserves the task’s label.

Define what a new test case means before splitting. If deployment requires recognizing unseen products, keep all views of each product in one split; for unseen videos, group by video. Repeated views may be appropriate when the intended task is a new view of an already-known object, but then the evaluation answers that narrower question. Keep derived crops and augmented copies with their source image. Perceptual hashes help find some near-duplicates; metadata and manual inspection are also needed for related images that look different.

The architectures these ideas came from are traced in CNN Architectures Compared: LeNet to EfficientNet.

Exercises

1. Freeze order matters. Using an untrained ResNet-18 to isolate the freezing behavior, build the model two ways: freeze all parameters and then replace the head, versus replace the head and then freeze all parameters. Report the trainable parameter count for each and explain the difference.

You should get: a working setup and one with nothing trainable at all.

import torch.nn as nn, torchvision
def build(freeze_first):
    m = torchvision.models.resnet18(weights=None)
    if freeze_first:
        for p in m.parameters(): p.requires_grad = False
        m.fc = nn.Linear(m.fc.in_features, 10)      # new head is trainable
    else:
        m.fc = nn.Linear(m.fc.in_features, 10)
        for p in m.parameters(): p.requires_grad = False   # freezes the head too
    return sum(p.numel() for p in m.parameters() if p.requires_grad)
print(build(True), build(False))
# 5130 0

2. Parameters and BatchNorm statistics. Set requires_grad=False on a BatchNorm layer, run several forward passes in training mode, and check whether its running_mean changed. Then repeat with .eval() on that module.

You should get: statistics that move despite the parameters being frozen, and statistics that hold still once eval mode is set.

import torch, torch.nn as nn
bn = nn.BatchNorm1d(4)
for p in bn.parameters(): p.requires_grad = False
bn.train()
before = bn.running_mean.clone()
for _ in range(5): bn(torch.randn(32, 4) + 10.0)
print("train mode moved:", not torch.allclose(before, bn.running_mean))
# train mode moved: True

bn.eval(); before = bn.running_mean.clone()
for _ in range(5): bn(torch.randn(32, 4) + 10.0)
print("eval mode moved:", not torch.allclose(before, bn.running_mean))
# eval mode moved: False

3. Which label should a transform preserve? Decide whether each is valid and why: horizontal flip on chest X-rays; vertical flip on satellite imagery; strong color jitter on a flower-species classifier; random crop at scale 0.08 on a fine-grained bird dataset.

For each case, identify the label, the visual evidence it needs, and any orientation information that must be preserved. The image category alone is not enough to settle every case.

References

• Torchvision. ResNet-50 weights and preprocessing; PyTorch. Transfer Learning for Computer Vision.
• Yosinski, Clune, Bengio, and Lipson (2014). How Transferable Are Features in Deep Neural Networks? NIPS.
• Zhang, Cisse, Dauphin, and Lopez-Paz (2018). mixup: Beyond Empirical Risk Minimization. ICLR.
• Yun et al. (2019). CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features. ICCV.
• Kornblith, Shlens, and Le (2019). Do Better ImageNet Models Transfer Better? CVPR.
• Cubuk, Zoph, Shlens, and Le (2020). RandAugment: Practical Automated Data Augmentation with a Reduced Search Space. CVPR Workshops.
• Zhong, Zheng, Kang, Li, and Yang (2020). Random Erasing Data Augmentation. AAAI.