Fine-Tune ResNet on Custom Data

You've just trained a model from scratch, waited hours, and still got 78% accuracy. Your dataset is tiny but your deadline isn't. Sound familiar? Training a deep neural network from random weights is expensive, slow, and data-hungry — the exact opposite of what you need when building real applications. Fine-tuning ResNet on custom data solves this by borrowing knowledge from a model already trained on millions of images, then adapting it to your specific problem in minutes, not days. In this lesson, you'll learn the mental model behind transfer learning, walk through a complete fine-tuning pipeline step by step, and get hands-on with code you can adapt to your own dataset.

The problem this lesson solves

If you've tried to train an image classifier from scratch, you've felt the pain:

• Not enough data — deep networks overfit badly on small datasets (a few thousand images is rarely enough).
• Not enough compute — training ResNet-50 from scratch on ImageNet takes days on multiple GPUs; you likely have one modest GPU (or a free Colab session).
• Not enough time — hyperparameter tuning, debugging, and re-training eat your sprint budget.

You need a way to build an accurate image classifier for your use case — be it product defects, plant diseases, or cat breeds — without starting from zero.

Pro tip: Many real-world AI projects fail not because the model is wrong, but because the training strategy is. Fine-tuning is the default, pragmatic answer for most business problems.

Core concept / mental model

Think of a pretrained network like a medical student — it has spent years learning general anatomy, imaging patterns, and diagnostic heuristics. It's not ready to specialize yet, but it doesn't need to relearn what a human looks like. Fine-tuning is the final semester of residency: you take that well-trained brain and drill it on your specific cases (your custom images).

Transfer learning is the umbrella term. Feature extraction (freezing the base) and fine-tuning (unfreezing and updating some or all weights) are two common strategies. With ResNet, you usually keep the convolutional backbone — which learns generic features like edges, textures, and shapes — and replace the final fully connected layer with one that matches your number of classes. Then you train, initially only the new head, and optionally later the whole network with a lower learning rate.

Why ResNet specifically? Its residual connections (skip connections) let gradients flow through many layers without vanishing, making it a reliable workhorse. Pretrained weights (ImageNet) give you a flying start.

How it works step by step

Fine-tuning ResNet on custom data follows a clear sequence:

1. Prepare your data — organize images into train/val/test folders, one subfolder per class. Ensure a balanced-ish number of images per class, and resize images to 224×224 (ResNet's expected input).

2. Load the pretrained model — In PyTorch, torchvision.models.resnet50(pretrained=True) loads weights trained on ImageNet.

3. Replace the classifier head — ResNet's fc layer outputs 1000 ImageNet classes. Swap it for a nn.Linear(2048, num_classes) (for ResNet-50) plus an appropriate activation if needed.

4. Freeze earlier layers — Prevent them from changing in the first phase (optional, but reduces compute and risk of overfitting).

5. Train the new head — Use a moderate learning rate (e.g., 1e-3) and standard optimizers like Adam or SGD with momentum.

6. Optionally unfreeze and fine-tune — After the head converges, unfreeze the top few layers (or all) and train with a much lower learning rate (e.g., 1e-5) to adapt features to your domain.

7. Evaluate — Check accuracy on a held-out test set. If overfitting, add data augmentation or dropout.

Hands-on walkthrough

Let's implement a complete fine-tuning pipeline with PyTorch and torchvision. We'll assume you have train/ and val/ folders with subfolders per class (e.g., cats/, dogs/).

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models, transforms

data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
}

image_datasets = {
    'train': datasets.ImageFolder('data/train', data_transforms['train']),
    'val': datasets.ImageFolder('data/val', data_transforms['val'])
}

dataloaders = {
    'train': torch.utils.data.DataLoader(image_datasets['train'], batch_size=32, shuffle=True),
    'val': torch.utils.data.DataLoader(image_datasets['val'], batch_size=32, shuffle=False)
}

# Load pretrained ResNet50
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
num_classes = len(image_datasets['train'].classes)
model.fc = nn.Linear(model.fc.in_features, num_classes)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# Freeze all layers except the new fc layer
for param in model.parameters():
    param.requires_grad = False
for param in model.fc.parameters():
    param.requires_grad = True

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=0.001)

Now train the classifier head for a few epochs:

from copy import deepcopy

def train_model(model, dataloaders, criterion, optimizer, num_epochs=5):
    best_model_wts = deepcopy(model.state_dict())
    best_acc = 0.0
    for epoch in range(num_epochs):
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()
            else:
                model.eval()
            running_loss = 0.0
            running_corrects = 0
            for inputs, labels in dataloaders[phase]:
                inputs, labels = inputs.to(device), labels.to(device)
                optimizer.zero_grad()
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)
            print(f'Epoch {epoch+1}/{num_epochs} - {phase}: Loss {epoch_loss:.4f} Acc {epoch_acc:.4f}')
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = deepcopy(model.state_dict())
    print(f'Best val Acc: {best_acc:4f}')
    model.load_state_dict(best_model_wts)
    return model

model = train_model(model, dataloaders, criterion, optimizer, num_epochs=5)

Expected output:

Epoch 1/5 - train: Loss 1.2310 Acc 0.4821
Epoch 1/5 - val: Loss 1.1023 Acc 0.5200
...
Best val Acc: 0.843750

Now unfreeze a few top layers and fine-tune the whole network with a low learning rate:

# Unfreeze all layers, but keep a low learning rate
for param in model.parameters():
    param.requires_grad = True
optimizer = optim.SGD(model.parameters(), lr=0.0001, momentum=0.9)

model = train_model(model, dataloaders, criterion, optimizer, num_epochs=5)

Expected improvement: Accuracy typically jumps a few points because earlier features adapt to your domain.

Pro tip: Save your best model with torch.save(model.state_dict(), 'resnet_finetuned.pth'). You'll need it for the next lesson.

Compare options / when to choose what

Not every situation calls for full fine-tuning. Here's a quick guide:

Simpler alternatives: - A smaller custom CNN if your dataset is tiny and simple, but expect lower accuracy. - Use a smaller ResNet variant (e.g., ResNet-18) for faster training with less compute, or a larger one (ResNet-101) when you have a big dataset and GPU budget. - Dropout regularization in the classifier head can prevent overfitting when fine-tuning on small datasets.

Troubleshooting & edge cases

• Accuracy stuck at low values (e.g., 50% with 2 classes) — Check that your images aren't mislabeled, and that your normalization matches ImageNet's statistics (if using pretrained weights). Also, verify the learning rate isn't too high — try lowering it.

• Exploding loss / NaN — This often happens with a learning rate that's too high, especially with batch size 1. Reduce the LR (e.g., from 0.001 to 0.0001) and add gradient clipping.

• Overfitting quickly — If training accuracy is high but validation is low, add data augmentation (random crops, flips, color jitter) or increase dropout in the head. Freeze more layers.

• Input size mismatch — ResNet expects (3, 224, 224) images. If your images are grayscale, convert to RGB (transforms.Grayscale(num_output_channels=3)).

• Class imbalance — Use weighted random sampler or weighted_cross_entropy to handle skewed classes.

• CUDA out of memory — Reduce batch size, use gradient accumulation, or switch to CPU (slowly).

What you learned & what's next

You've learned the core idea behind fine-tuning ResNet on custom data — how to leverage pretrained features, replace the classifier head, freeze layers, and train with a low learning rate. You can now apply this to your own image classification problem. In the next lesson, we'll take this finetuned model and deploy it as an inference service, so you can turn your experiment into a real API.