Debugging a Fine-Tuning Script: ResNet-50 on Three Images

Going Deeper explained how ResNet works. This post uses one. It takes a short PyTorch notebook that is meant to fine-tune an ImageNet-pretrained ResNet-50 on a tiny data set of cats, pineapples and swans, and that does not work. The notebook is shown exactly as written; the post reads it line by line, reproduces what it does, and then fixes it one problem at a time, measuring the effect of every change.

The script is short, and almost every line hides a lesson: what the loss function expects, what zero_grad is for, why a pretrained network needs its normalisation, how BatchNorm behaves with one image per batch, and why PyTorch’s version matters. All experiments were run on a laptop CPU with fixed seeds; the scripts are linked at the end.

---

1. The Task and the Data

The folder contains a notebook, a list of the 1,000 ImageNet classes, and a toy data set:

Finetune ResNet50 Task/
├── Finetune ResNet50 Classification.ipynb
├── imagenet_classes.txt          # "0, tench" … "999, toilet_tissue"
└── DummyData/
    ├── train/
    │   ├── cat/cat.png
    │   ├── pineapple/pineapple.png
    │   └── swan/black_swan_1.png
    └── test/
        ├── cat/cat.png
        ├── pineapple/pineapple.png
        └── swan/black_swan_2.png

One 1024 × 1024 image per class for training, a different one per class for testing. All three classes already exist in ImageNet: tiger cat (id 282), pineapple (953) and black swan (100). The notebook also expects the pretrained weights in pretrained_models/resnet50-19c8e357.pth, which is not included.

PyTorch’s ImageFolder assigns labels to the class folders in alphabetical order: cat = 0, pineapple = 1, swan = 2. The notebook’s helper get_imagenet_class maps 0 → 282, 1 → 953 and anything else → 100, which matches that order. The catch, shown in the diagram, is where the mapping is used.

---

2. The Code, as Written

The notebook has one code cell:

import os
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision.models import resnet50
 
def get_imagenet_class(cid):
    # hard coded to map our simple dataset to imagenet numbers
    return 282 if cid == 0 else 953 if cid == 1 else 100
 
# Load the class labels for ImageNet
with open('imagenet_classes.txt') as f:
    class_names = [line.split(",")[1].strip() for line in f.readlines()]
    
# Set the device to GPU if available, otherwise use CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
# Define data directories for training and testing
train_dir = './DummyData/train'
test_dir = './DummyData/test'
 
# Define data transforms
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.1, 0.1, 0.1]),
])
 
# Create custom ImageNet dataset loaders
train_dataset = ImageFolder(train_dir, transform=transform)
test_dataset = ImageFolder(test_dir, transform=transform)
 
# Create data loaders
train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=True, num_workers=4)
 
# Load a pre-trained ResNet-18 model
model = resnet50(pretrained=False)
# Make sure to download the pretrained model from https://download.pytorch.org/models/resnet50-19c8e357.pth
model.load_state_dict(torch.load('pretrained_models/resnet50-19c8e357.pth'))
 
# Move the model to the selected device (CPU or GPU)
model.to(device)
 
# Define the loss function and optimizer
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.9, momentum=0.001)
 
# Training loop
num_epochs = 3
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {running_loss / len(train_loader)}")
 
# Evaluate the model on the test set
model.eval()
 
with torch.no_grad():
    for inputs, labels in test_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        in_label = get_imagenet_class(labels.item())
        print(f"True Class: {class_names[in_label]}, Predicted Class: {class_names[predicted.item()]}")

When it was last run, the saved output ended with a warning and an error:

UserWarning: Using a target size (torch.Size([1])) that is different to the input size
(torch.Size([1, 1000])). This will likely lead to incorrect results due to broadcasting.
Please ensure they have the same size.
  return F.mse_loss(input, target, reduction=self.reduction)
 
RuntimeError: Found dtype Long but expected Float

---

3. Reproducing It

Two things are needed besides the folder: the pretrained weights (from the URL in the code’s comment) and a PyTorch version. The version matters more than usual, because PyTorch has changed its behaviour twice on exactly the lines this notebook uses:

cd "Finetune ResNet50 Task"
mkdir -p pretrained_models
curl -L -o pretrained_models/resnet50-19c8e357.pth https://download.pytorch.org/models/resnet50-19c8e357.pth
 
python3.12 -m venv .venv && source .venv/bin/activate
pip install "torch==2.5.1" "torchvision==0.20.1" "numpy<2" jupyter
jupyter nbconvert --to notebook --execute --allow-errors "Finetune ResNet50 Classification.ipynb" --output run.ipynb

Run unchanged, the notebook does different things depending on the version:

PyTorch	What happens when the unchanged notebook runs	Why
2.5.1 (and earlier)	The warning above, then RuntimeError: Found dtype Long but expected Float in the first loss.backward(): the saved output, reproduced exactly	MSELoss’s backward pass needs a floating-point target; the labels are integers
2.6 and later (tested: 2.14.1)	Stops earlier, at torch.load: “Cannot use weights_only=True with files saved in the legacy .tar format”	Since 2.6, torch.load defaults to weights_only=True, and this 2017 weight file uses the old format
2.14.1, once the weights load	No error at all. Three epochs run, and every test image is then classified as the same unrelated class (skunk with seed 0)	MSELoss now accepts the integer target, broadcasts it, and trains on a meaningless objective

The last row is the most instructive: the bug that used to stop the script now lets it run to the end and print confident nonsense. A crash is the friendly version of a bug.

(For the experiments below, the weights are loaded with weights_only=False, which is safe here because the file comes from PyTorch’s own server; for files from unknown sources it is not.)

---

4. Reading It Line by Line

Imports. Standard PyTorch and torchvision; torchvision.transforms is imported twice, which is harmless.

get_imagenet_class. Maps the data set’s labels to ImageNet ids, used only in the evaluation, to print the true class by name.

class_names. Reads imagenet_classes.txt (“0, tench”, …) and keeps the name after the comma: class_names[282] is tiger_cat.

device. The first GPU if there is one, otherwise the CPU.

Transforms.

• Resize((224, 224)) scales every image to 224 × 224, the input size ResNet was trained at (for non-square images this would distort them; these are square).
• ToTensor() turns the image into a 3 × 224 × 224 tensor with values in [0, 1].
• Normalize(mean=[0.5]*3, std=[0.1]*3) computes $(x-0.5)/0.1$ per channel, giving values from −5 to 5.

Datasets and loaders. ImageFolder reads the class folders (labels 0, 1, 2); DataLoader serves one image at a time (batch_size=1), shuffled, with 4 worker processes.

Model. resnet50(pretrained=False) builds the architecture with random weights (the pretrained argument is deprecated since torchvision 0.13), then load_state_dict(torch.load(...)) loads the ImageNet weights. (The comment says ResNet-18; the code loads ResNet-50.) The final layer model.fc maps 2,048 features to 1,000 ImageNet classes.

Loss and optimiser. nn.MSELoss(), the squared error of Linear Regression, and stochastic gradient descent with learning rate 0.9 and momentum 0.001.

Training loop. Three epochs; for each image: forward pass, loss against the label, loss.backward(), optimizer.step(). There is no optimizer.zero_grad().

Evaluation. model.eval(), no gradients, and for each test image the arg-max of the 1,000 outputs, printed by name next to the mapped true class.

---

5. What Goes Wrong, One Fix at a Time

Applying the fixes cumulatively, with three seeds each (seed 0 shown in the lab):

Setup (PyTorch 2.14.1, CPU)	Test images correct, seeds 0 · 1 · 2	What the model predicts (seed 0)
Pretrained, no training, notebook normalisation	0 (no training)	spotlight 0.74, perfume 0.40, torch 0.98
Pretrained, no training, ImageNet normalisation	3 (no training)	tiger cat 0.48, pineapple 1.00, black swan 0.88
A. The notebook as written	0 · 0 · 0	skunk, for all three
B. + cross-entropy, labels mapped to ImageNet ids	1 · 1 · 1	black swan, for all three
C. + zero_grad, + ImageNet normalisation	1 · 1 · 1	pineapple, for all three
D. + lr 0.001, momentum 0.9	3 · 3 · 3	tiger cat 0.97, pineapple 1.00, black swan 0.99
E. D with BatchNorm frozen	0 · 1 · 3	NaN: the loss diverged
F. E with lr 0.0001	3 · 3 · 3	tiger cat 0.92, pineapple 1.00, black swan 0.99
G. New 3-class head, backbone frozen with requires_grad=False	1 · 2 · 1	pineapple, for all three
H. G with BatchNorm also frozen	1 · 3 · 3	correct on two of three seeds

5.1 The loss: MSELoss on a class label

nn.MSELoss() expects a target with the same shape as the output, here 1 × 1,000. The notebook gives it the label, a tensor of shape 1 holding 0, 1 or 2. PyTorch broadcasts it to all 1,000 outputs, which is what the warning says, and the loss becomes

$$\frac{1}{1000}\sum _{k=1}^{1000}{\left(z_k-y\right)}^2,$$

the average squared distance of every output from the number $y$. Training pushes all 1,000 outputs toward 0 for a cat, toward 1 for a pineapple and toward 2 for a swan, regardless of which class each output stands for; nothing in this loss says which output should be the largest. On old PyTorch the backward pass refuses the integer target; on new PyTorch it runs, and the result (skunk for everything) is whatever class happens to come out on top.

A classification loss treats the label as an index. nn.CrossEntropyLoss() applies a softmax to the outputs and maximises the log-probability of the one class named by the label (the cross-entropy of Classification):

5.2 The label space: 0, 1, 2 against 1,000 classes

Even with cross-entropy, the label must refer to the right output. The network has 1,000 outputs for the ImageNet classes; label 0 means output 0, which is tench. The notebook maps the labels with get_imagenet_class in the evaluation loop, but not in the training loop. There are two consistent choices (diagram in §6): map the labels to 282, 953 and 100 during training as well, or replace the 1,000-way head with a new 3-way layer and keep the labels 0, 1, 2.

5.3 The missing zero_grad

loss.backward() adds each parameter’s gradient to its .grad field; it never resets it. PyTorch does this on purpose (it allows accumulating gradients over several small batches), so every training step must start with optimizer.zero_grad(). Without it, the update at step $t$ uses $g_1+g_2+\cdots +g_t$, which grows with every step.

5.4 The learning rate and momentum

lr=0.9 with momentum=0.001 looks like the two numbers were swapped. Fine-tuning a pretrained network usually uses a learning rate around 0.001 to 0.0001 for SGD, with momentum about 0.9. At 0.9, the first update changes the weights so much that the network collapses: in B and C, every test image was assigned to the same class. (Momentum 0.001 is effectively none.) With lr 0.001 and momentum 0.9 (D), fine-tuning works on every seed.

5.5 The normalisation

A pretrained network has only ever seen inputs normalised with ImageNet’s statistics, per colour channel: mean (0.485, 0.456, 0.406) and standard deviation (0.229, 0.224, 0.225). The notebook’s mean 0.5 and standard deviation 0.1 make the inputs about 2.3 times larger: the test images span −5.0 to 5.0 (standard deviation 2.74) instead of −2.1 to 2.6 (1.17).

The measurement that matters most in this post: without any training, the pretrained ResNet-50 classifies all three test images correctly with ImageNet normalisation, and none with the notebook’s (spotlight, perfume, torch). Every result above with the notebook’s normalisation started from a broken network.

5.6 BatchNorm with one image per batch

ResNet-50 has 53 batch-normalisation layers (Ioffe & Szegedy, 2015). In training mode, each normalises its input by the mean and variance of the current batch, and updates a running average that is used in evaluation mode. With batch_size=1, the “batch” statistics come from a single image. This had three measurable consequences:

1. In training mode, with a small learning rate, it worked (D: 3 of 3 on every seed). Per-image normalisation keeps the activations in a sane range.
2. Freezing BatchNorm (evaluation mode) at the same learning rate diverged on two of three seeds (E). Without the per-image normalisation, the gradients from single images were large (a gradient norm of about 100 in the first steps), and momentum carried them into NaN. A ten times smaller learning rate (F) was stable.
3. requires_grad=False does not freeze BatchNorm. Freezing the backbone’s weights and training only a new head (G) failed, because in training mode the frozen network’s BatchNorm layers kept overwriting their ImageNet running statistics with statistics from single images. Only after also putting them in evaluation mode (H) did the head train properly.

5.7 Loading the weights

resnet50(pretrained=False) is deprecated; the current way is resnet50(weights=None) followed by loading, or simply resnet50(weights=ResNet50_Weights.IMAGENET1K_V1), which downloads ImageNet weights in the current format. torch.load also needs map_location="cpu" to load GPU-saved weights on a machine without a GPU, and, for this legacy file on PyTorch 2.6 or later, weights_only=False.

5.8 Smaller things

num_workers=4 starts four worker processes to load three images, and in a notebook on Windows or macOS this can hang; num_workers=0 is the safe choice for small data. Shuffling the test set changes only the order of the printed lines. Resize((224, 224)) ignores the aspect ratio, which is fine for these square images; the standard ImageNet evaluation is Resize(256) followed by CenterCrop(224).

---

6. Does This Even Need Fine-Tuning?

Both strategies work once the bugs are gone (D and F for A; H for B). But the data set is far too small for either to mean much:

• The pretrained model already solves the task. All three classes are ImageNet classes, and with the right normalisation the untrained model gets 3 of 3. Fine-tuning can only make that worse or keep it.
• Three training images cannot train 25.6 million parameters (strategy A) or even 6,147 (the new head in B). In H, a head that reached near-zero training loss still misclassified two of three test images on one seed: with one example per class, which direction in the 2,048-dimensional feature space separates the classes is barely constrained.
• Three test images cannot measure anything. One wrong image changes the accuracy by 33 points, and the “cat” target is tiger cat specifically: the pretrained model’s 0.48 for tiger cat against 0.32 for tabby means another cat photo could easily be scored “wrong” for being a different cat.

A realistic fine-tuning setup uses tens to thousands of images per class, a validation set to choose the learning rate and the number of epochs (Cross-Validation and the Bootstrap), data augmentation, and usually trains the new head first with the backbone frozen before unfreezing it with a smaller learning rate (Kumar et al., 2022). How well pretrained features transfer depends on how close the new task is to ImageNet (Yosinski et al., 2014; Kornblith et al., 2019): here it is ImageNet.

---

7. A Corrected Version

The original notebook stays as it is above. For comparison, a version with the fixes applied (strategy A, configuration F of §5):

import torch, torch.nn as nn, torch.optim as optim
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from torchvision.models import resnet50
 
IMAGENET_ID = {0: 282, 1: 953, 2: 100}           # ImageFolder labels (alphabetical) -> ImageNet ids
with open("imagenet_classes.txt") as f:
    class_names = [line.split(",")[1].strip() for line in f if line.strip()]
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),   # ImageNet statistics
])
train_loader = DataLoader(ImageFolder("./DummyData/train", transform), batch_size=1, shuffle=True, num_workers=0)
test_loader = DataLoader(ImageFolder("./DummyData/test", transform), batch_size=1, shuffle=False, num_workers=0)
 
model = resnet50(weights=None)
model.load_state_dict(torch.load("pretrained_models/resnet50-19c8e357.pth", map_location="cpu", weights_only=False))
model.to(device)
 
criterion = nn.CrossEntropyLoss()                                    # labels are class indices
optimizer = optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
 
for epoch in range(3):
    model.train()
    for m in model.modules():                                        # keep the ImageNet BatchNorm statistics
        if isinstance(m, nn.BatchNorm2d):
            m.eval()
    running_loss = 0.0
    for inputs, labels in train_loader:
        targets = torch.tensor([IMAGENET_ID[l] for l in labels.tolist()])
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()                                        # gradients do not accumulate
        loss = criterion(model(inputs), targets)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f"Epoch {epoch + 1}/3, Loss: {running_loss / len(train_loader):.4f}")
 
model.eval()
with torch.no_grad():
    for inputs, labels in test_loader:
        predicted = model(inputs.to(device)).argmax(1).item()
        print(f"True: {class_names[IMAGENET_ID[labels.item()]]}, predicted: {class_names[predicted]}")

Run three times (without a fixed seed), it printed the correct class for all three test images every time, with the training loss falling from about 0.19 to 0.001–0.002 in three epochs. For three classes that are already in ImageNet, the honest final answer is shorter still: skip the training loop and evaluate the pretrained model with the right normalisation.

---

8. Common Claims, Checked

Claim	Verdict
”MSELoss works for classification if the labels are numbers.”	It treats the label as a value to approach: here it pulled all 1,000 outputs toward 0, 1 or 2, and every test image became skunk. Use cross-entropy with class indices.
”If the script runs without errors, the training is fine.”	On PyTorch 2.14 the notebook runs to the end without an error and gets 0 of 3; the same code crashes on 2.5.1. Check the outputs, not just the exit code.
”The labels from ImageFolder can be fed to the pretrained 1,000-class head.”	They are 0, 1, 2 (tench, goldfish, great white shark in ImageNet). Map them to 282, 953, 100, or replace the head.
”zero_grad() is optional.”	backward() adds to .grad; without zero_grad(), each update uses the sum of all previous gradients.
”A higher learning rate fine-tunes faster.”	At lr 0.9 every configuration collapsed onto one class; 0.001 (BatchNorm in training mode) or 0.0001 (frozen BatchNorm) worked on all seeds.
”Any normalisation works; the network adapts.”	The pretrained model went from 3 of 3 (ImageNet statistics) to 0 of 3 (mean 0.5, std 0.1) without any training. A pretrained network expects its own normalisation.
”requires_grad=False freezes a layer.”	It freezes the weights. BatchNorm still updates its running statistics in training mode; the frozen-backbone head scored 1 · 2 · 1 of 3 until BatchNorm was put in evaluation mode too (1 · 3 · 3).
”Freezing BatchNorm always makes fine-tuning more stable.”	With batch size 1 at lr 0.001 it diverged to NaN on two of three seeds; it needed lr 0.0001.
”torch.load on an official checkpoint always works.”	Since PyTorch 2.6 this 2017 file fails with the new weights_only=True default; use the torchvision weights API, or weights_only=False for trusted files only.
”Fine-tuning is needed to classify cats, pineapples and swans.”	They are ImageNet classes: the pretrained model got all three right without training. Three images could not show that fine-tuning helps anyway.

---

Summary

Problem in the notebook	Fix
nn.MSELoss() on class labels	nn.CrossEntropyLoss()
Labels 0–2 against 1,000 ImageNet outputs	Map to 282, 953, 100 in training too, or a new 3-way head
No optimizer.zero_grad()	Call it before every backward()
lr=0.9, momentum=0.001	lr ≈ 1e-3 to 1e-4, momentum=0.9
Normalize(0.5, 0.1)	ImageNet mean and standard deviation
BatchNorm with batch size 1	Small learning rate; freeze BatchNorm (eval mode) when freezing the backbone
pretrained=False + torch.load of a legacy file	weights=ResNet50_Weights.IMAGENET1K_V1, or map_location and weights_only=False for trusted files
One image per class	More data, a validation set, and first check whether the pretrained model already does the job

Next: Efficient CNNs: Doing More With Less: 1×1 convolutions, Inception, MobileNet and EfficientNet, or how to get more accuracy per parameter and per multiply.

The experiments: debug_finetune.py (PyTorch 2.14.1, torchvision 0.29.1, CPU, seeds 0–2), run in the task folder with the official weights; the unchanged notebook was executed with jupyter nbconvert under PyTorch 2.5.1 and 2.14.1. The three images are not reproduced here.

---

References

Architecture and data

• He, K., Zhang, X., Ren, S., & Sun, J. (2016). Deep residual learning for image recognition. CVPR.
• Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., & Fei-Fei, L. (2009). ImageNet: A large-scale hierarchical image database. CVPR.
• Ioffe, S., & Szegedy, C. (2015). Batch normalization: Accelerating deep network training by reducing internal covariate shift. ICML.

Transfer learning and fine-tuning

• Yosinski, J., Clune, J., Bengio, Y., & Lipson, H. (2014). How transferable are features in deep neural networks? NeurIPS.
• Kornblith, S., Shlens, J., & Le, Q. V. (2019). Do better ImageNet models transfer better? CVPR.
• Kumar, A., Raghunathan, A., Jones, R., Ma, T., & Liang, P. (2022). Fine-tuning can distort pretrained features and underperform out-of-distribution. ICLR.

Tools

• Paszke, A., Gross, S., Massa, F., et al. (2019). PyTorch: An imperative style, high-performance deep learning library. NeurIPS.
• PyTorch documentation: torch.load (the weights_only default since 2.6), nn.CrossEntropyLoss, nn.MSELoss, and torchvision’s ResNet-50 weights.