Image Classification with Transfer Learning: A Comprehensive Tutorial

Table of Contents

Introduction

Transfer learning is a machine learning technique where a model trained on a large dataset is repurposed for a different but related task. In computer vision, this typically involves taking a model pre-trained on ImageNet (1.4 million images, 1000 classes) and adapting it to your specific classification problem. This approach dramatically reduces the data and compute needed to achieve high accuracy.

This tutorial covers the complete workflow: selecting a pre-trained model, preparing your data, fine-tuning with PyTorch, evaluating performance, and deploying the final model.

Prerequisites

pip install torch torchvision
pip install timm                 # PyTorch Image Models - extensive model zoo
pip install albumentations       # Advanced augmentation
pip install scikit-learn         # Metrics
pip install matplotlib seaborn   # Visualization
pip install onnx onnxruntime     # Export and deployment

• Python 3.8 or higher
• GPU with at least 6 GB VRAM (training), CPU is sufficient for inference
• Basic understanding of neural networks and PyTorch

import torch
import torchvision
import timm

print(f"PyTorch: {torch.version}")
print(f"Torchvision: {torchvision.version}")
print(f"CUDA: {torch.cuda.isavailable()}")

print(f"Timm: {timm.version}")
print(f"Available timm models: {len(timm.listmodels())}")

Understanding Transfer Learning

Transfer learning works because the early layers of a CNN learn universal features (edges, textures, patterns) that apply to almost any vision task. Only the later layers become task-specific.

There are two main strategies:

# Strategy 1: Feature extraction
def createfeatureextractor(modelname, numclasses):
    """Freeze all layers except the classification head."""
    model = timm.createmodel(modelname, pretrained=True, numclasses=numclasses)

    # Freeze all parameters
    for param in model.parameters():
        param.requiresgrad = False


    # Unfreeze the classification head
    if hasattr(model, 'classifier'):
        for param in model.classifier.parameters():
            param.requiresgrad = True
    elif hasattr(model, 'fc'):
        for param in model.fc.parameters():
            param.requiresgrad = True

    elif hasattr(model, 'head'):
        for param in model.head.parameters():
            param.requiresgrad = True

    return model

Strategy 2: Full fine-tuning with discriminative learning rates
def createfinetuningmodel(modelname, numclasses):
    """All layers trainable, but with different learning rates."""
    model = timm.createmodel(modelname, pretrained=True, numclasses=numclasses)
    # All parameters are trainable by default
    return model

Choosing a Pre-trained Model

Comparing Popular Architectures

import timm

def comparemodels():

    """Compare popular models by size, speed, and accuracy."""
    modelsinfo = {
        "resnet50": {
            "params": "25.6M",
            "top1imagenet": 80.4,

            "speed": "fast",
            "usecase": "General purpose, good baseline"
        },
        "efficientnetb3": {

            "params": "12.2M",
            "top1imagenet": 82.0,
            "speed": "medium",
            "usecase": "Best accuracy-efficiency tradeoff"

        },
        "efficientnetv2s": {
            "params": "21.5M",
            "top1imagenet": 83.9,

            "speed": "medium",
            "usecase": "Improved EfficientNet, faster training"
        },
        "vitbasepatch16224": {

            "params": "86.6M",
            "top1imagenet": 84.5,
            "speed": "slow",
            "usecase": "Vision Transformer, best with large datasets"

        },
        "convnexttiny": {
            "params": "28.6M",
            "top1imagenet": 82.1,

            "speed": "medium",
            "usecase": "Modern CNN matching ViT performance"
        },
    }

    for name, info in modelsinfo.items():

        print(f"{name:30s} | Params: {info['params']:>8s} | "
              f"Top-1: {info['top1imagenet']:.1f}% | "
              f"Speed: {info['speed']:>6s} | {info['usecase']}")


    return modelsinfo

comparemodels()

import timm

def loadpretrainedmodel(modelname, numclasses, pretrained=True):

    """
    Load a pre-trained model and adapt it for your number of classes.
    timm automatically replaces the classification head.
    """
    model = timm.createmodel(
        modelname,

        pretrained=pretrained,
        numclasses=numclasses,

    )

    # Get model-specific preprocessing config
    dataconfig = timm.data.resolvemodeldataconfig(model)

    print(f"Model: {modelname}")
    print(f"Input size: {dataconfig['inputsize']}")
    print(f"Mean: {dataconfig['mean']}")

    print(f"Std: {dataconfig['std']}")
    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

    return model, dataconfig


model, config = loadpretrainedmodel("efficientnetb3", numclasses=10)

dataset/
    train/
        classa/
            img001.jpg
            img002.jpg
        classb/

            img003.jpg
            img004.jpg
    val/
        classa/
            img005.jpg
        classb/

            img006.jpg

import torch
from torch.utils.data import DataLoader
from torchvision import datasets
import albumentations as A
from albumentations.pytorch import ToTensorV2
import numpy as np
from PIL import Image

class AlbumentationsDataset(torch.utils.data.Dataset):
    """Wraps a torchvision dataset with Albumentations transforms."""

    def init(self, rootdir, transform=None):
        self.dataset = datasets.ImageFolder(rootdir)

        self.transform = transform
        self.classes = self.dataset.classes

    def len(self):
        return len(self.dataset)

    def getitem(self, idx):
        image, label = self.dataset[idx]
        image = np.array(image)

        if self.transform:
            augmented = self.transform(image=image)
            image = augmented["image"]

        return image, label

def createdataloaders(datadir, inputsize=224, batchsize=32, numworkers=4):

    """Create train and validation data loaders with proper augmentation."""

    traintransform = A.Compose([
        A.RandomResizedCrop(height=inputsize, width=inputsize, scale=(0.8, 1.0)),
        A.HorizontalFlip(p=0.5),
        A.VerticalFlip(p=0.1),
        A.Rotate(limit=15, p=0.3),
        A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=0.5),
        A.GaussNoise(varlimit=(10, 50), p=0.2),

        A.CoarseDropout(maxholes=8, maxheight=16, maxwidth=16, p=0.2),
        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        ToTensorV2(),
    ])

    valtransform = A.Compose([

        A.Resize(height=int(inputsize  1.14), width=int(inputsize  1.14)),

        A.CenterCrop(height=inputsize, width=inputsize),

        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        ToTensorV2(),
    ])

    traindataset = AlbumentationsDataset(f"{datadir}/train", transform=traintransform)
    valdataset = AlbumentationsDataset(f"{datadir}/val", transform=valtransform)


    trainloader = DataLoader(
        traindataset, batchsize=batchsize, shuffle=True,

        numworkers=numworkers, pinmemory=True, droplast=True,

    )
    valloader = DataLoader(
        valdataset, batchsize=batchsize, shuffle=False,

        numworkers=numworkers, pinmemory=True,
    )

    print(f"Training samples: {len(traindataset)}")

    print(f"Validation samples: {len(valdataset)}")
    print(f"Classes: {traindataset.classes}")


    return trainloader, valloader, traindataset.classes

trainloader, valloader, classes = createdataloaders("./dataset", inputsize=300, batchsize=32)

Fine-tuning with PyTorch

Setting Up the Training Configuration

import torch.nn as nn
import torch.optim as optim
from torch.optim.lrscheduler import CosineAnnealingLR


def setuptraining(model, numepochs=30, lr=1e-4, weightdecay=1e-4):
    """Configure optimizer, loss function, and scheduler."""

    # Discriminative learning rates: lower for pre-trained layers,
    # higher for the new classification head
    paramgroups = []

    headparams = []
    backboneparams = []


    for name, param in model.namedparameters():
        if not param.requiresgrad:

            continue
        if "classifier" in name or "fc" in name or "head" in name:
            headparams.append(param)
        else:
            backboneparams.append(param)


    paramgroups = [
        {"params": backboneparams, "lr": lr  0.1},   # 10x lower LR for backbone

        {"params": headparams, "lr": lr},               # Full LR for head
    ]

    optimizer = optim.AdamW(paramgroups, weightdecay=weightdecay)


    scheduler = CosineAnnealingLR(
        optimizer,
        Tmax=numepochs,

        etamin=lr  0.01,
    )

    criterion = nn.CrossEntropyLoss(labelsmoothing=0.1)


    return optimizer, scheduler, criterion

optimizer, scheduler, criterion = setuptraining(model, numepochs=30, lr=1e-4)

import time
from collections import defaultdict

class Trainer:
    """Complete training pipeline for image classification."""

    def init(self, model, optimizer, scheduler, criterion, device="cuda"):
        self.model = model.to(device)
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.criterion = criterion
        self.device = device
        self.history = defaultdict(list)
        self.bestvalacc = 0.0


    def trainoneepoch(self, trainloader):
        """Train for one epoch."""
        self.model.train()
        runningloss = 0.0

        correct = 0
        total = 0

        for batchidx, (images, labels) in enumerate(trainloader):

            images = images.to(self.device)
            labels = labels.to(self.device)

            # Forward pass
            outputs = self.model(images)
            loss = self.criterion(outputs, labels)

            # Backward pass
            self.optimizer.zerograd()
            loss.backward()

            # Gradient clipping
            torch.nn.utils.clipgradnorm(self.model.parameters(), maxnorm=1.0)

            self.optimizer.step()

            # Track metrics
            runningloss += loss.item()  images.size(0)

            , predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

        epochloss = runningloss / total
        epochacc = 100.0  correct / total
        return epochloss, epochacc


    @torch.nograd()
    def validate(self, valloader):

        """Evaluate on the validation set."""
        self.model.eval()
        runningloss = 0.0
        correct = 0
        total = 0

        for images, labels in valloader:

            images = images.to(self.device)
            labels = labels.to(self.device)

            outputs = self.model(images)
            loss = self.criterion(outputs, labels)

            runningloss += loss.item()  images.size(0)

            , predicted = outputs.max(1)

            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

        epochloss = runningloss / total

        epochacc = 100.0  correct / total
        return epochloss, epochacc

    def fit(self, trainloader, valloader, numepochs, savepath="bestmodel.pth"):

        """Complete training loop with validation and model saving."""
        print(f"{'Epoch':>6} | {'Train Loss':>10} | {'Train Acc':>9} | "
              f"{'Val Loss':>10} | {'Val Acc':>9} | {'LR':>10} | {'Time':>6}")
        print("-"  75)


        for epoch in range(numepochs):
            starttime = time.time()


            trainloss, trainacc = self.trainoneepoch(trainloader)
            valloss, valacc = self.validate(valloader)


            self.scheduler.step()
            currentlr = self.optimizer.paramgroups[0]["lr"]

            elapsed = time.time() - starttime

            # Log history
            self.history["trainloss"].append(trainloss)
            self.history["trainacc"].append(trainacc)
            self.history["valloss"].append(valloss)
            self.history["valacc"].append(valacc)

            print(f"{epoch + 1:>6} | {trainloss:>10.4f} | {trainacc:>8.2f}% | "
                  f"{valloss:>10.4f} | {valacc:>8.2f}% | {currentlr:>10.6f} | "

                  f"{elapsed:>5.1f}s")

            # Save best model
            if valacc > self.bestvalacc:
                self.bestvalacc = valacc

                torch.save({
                    "epoch": epoch,
                    "modelstatedict": self.model.statedict(),
                    "optimizerstatedict": self.optimizer.statedict(),

                    "valacc": valacc,

                    "valloss": valloss,

                }, savepath)
                print(f"       -> Saved best model (valacc: {valacc:.2f}%)")

        print(f"\nBest validation accuracy: {self.bestvalacc:.2f}%")
        return self.history

Run training
trainer = Trainer(model, optimizer, scheduler, criterion)
history = trainer.fit(trainloader, valloader, numepochs=30)

import matplotlib.pyplot as plt

def plottraininghistory(history):

    """Plot training and validation metrics."""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

    # Loss plot
    ax1.plot(history["trainloss"], label="Train Loss")
    ax1.plot(history["valloss"], label="Val Loss")

    ax1.setxlabel("Epoch")
    ax1.setylabel("Loss")

    ax1.settitle("Loss over Epochs")
    ax1.legend()
    ax1.grid(True)

    # Accuracy plot
    ax2.plot(history["trainacc"], label="Train Accuracy")

    ax2.plot(history["valacc"], label="Val Accuracy")
    ax2.setxlabel("Epoch")

    ax2.setylabel("Accuracy (%)")
    ax2.settitle("Accuracy over Epochs")

    ax2.legend()
    ax2.grid(True)

    plt.tightlayout()
    plt.savefig("traininghistory.png", dpi=150)

    plt.close()

plottraininghistory(history)

from sklearn.metrics import classificationreport, confusionmatrix
import seaborn as sns
import numpy as np

@torch.nograd()
def evaluatemodel(model, valloader, classnames, device="cuda"):

    """Generate comprehensive evaluation metrics."""
    model.eval()
    allpredictions = []
    alllabels = []

    allprobabilities = []

    for images, labels in valloader:

        images = images.to(device)
        outputs = model(images)
        probabilities = torch.softmax(outputs, dim=1)

        , predicted = outputs.max(1)
        allpredictions.extend(predicted.cpu().numpy())

        alllabels.extend(labels.numpy())
        allprobabilities.extend(probabilities.cpu().numpy())


    allpredictions = np.array(allpredictions)

    alllabels = np.array(alllabels)


    # Classification report
    print("Classification Report:")
    print(classificationreport(alllabels, allpredictions,
                                targetnames=classnames, digits=3))

    # Confusion matrix
    cm = confusionmatrix(alllabels, allpredictions)

    fig, ax = plt.subplots(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
                xticklabels=classnames, yticklabels=classnames, ax=ax)

    ax.setxlabel("Predicted")
    ax.setylabel("True")

    ax.settitle("Confusion Matrix")
    plt.tightlayout()

    plt.savefig("confusionmatrix.png", dpi=150)
    plt.close()

    accuracy = (allpredictions == alllabels).mean()  100
    print(f"\nOverall Accuracy: {accuracy:.2f}%")

    return accuracy

evaluatemodel(model, valloader, classes)

Model Export and Deployment

Export to ONNX

def exporttoonnx(model, inputsize=224, outputpath="model.onnx"):
    """Export the model to ONNX format for cross-platform deployment."""
    model.eval()
    model.cpu()

    dummyinput = torch.randn(1, 3, inputsize, inputsize)


    torch.onnx.export(
        model,
        dummyinput,
        outputpath,

        exportparams=True,
        opsetversion=17,

        doconstantfolding=True,

        inputnames=["input"],
        outputnames=["output"],

        dynamicaxes={
            "input": {0: "batchsize"},

            "output": {0: "batchsize"},
        },
    )
    print(f"Model exported to {outputpath}")


exporttoonnx(model, inputsize=300)

ONNX Runtime Inference

import onnxruntime as ort
import cv2
import numpy as np

class ONNXClassifier:
    """Run inference using ONNX Runtime."""

    def init(self, modelpath, classnames, inputsize=224):

        self.session = ort.InferenceSession(
            modelpath,
            providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
        )
        self.classnames = classnames
        self.inputsize = inputsize
        self.inputname = self.session.getinputs()[0].name

    def preprocess(self, imagepath):

        """Preprocess an image for inference."""
        image = cv2.imread(imagepath)
        image = cv2.cvtColor(image, cv2.COLORBGR2RGB)

        image = cv2.resize(image, (self.inputsize, self.inputsize))

        image = image.astype(np.float32) / 255.0

        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        image = (image - mean) / std

        image = image.transpose(2, 0, 1)  # HWC -> CHW
        image = np.expanddims(image, axis=0).astype(np.float32)
        return image

    def predict(self, imagepath, topk=5):
        """Run inference and return top-k predictions."""
        inputtensor = self.preprocess(imagepath)
        outputs = self.session.run(None, {self.inputname: inputtensor})
        logits = outputs[0][0]

        # Apply softmax
        explogits = np.exp(logits - np.max(logits))

        probabilities = explogits / explogits.sum()


        topindices = np.argsort(probabilities)[::-1][:topk]

        results = [
            {"class": self.classnames[i], "probability": float(probabilities[i])}
            for i in topindices

        ]

        return results

Usage
classifier = ONNXClassifier("model.onnx", classnames=classes, inputsize=300)

predictions = classifier.predict("testimage.jpg")
for pred in predictions:
    print(f"{pred['class']}: {pred['probability']:.4f}")

FastAPI Deployment

from fastapi import FastAPI, File, UploadFile
import io
from PIL import Image

app = FastAPI()
classifier = ONNXClassifier("model.onnx", classnames=classes, inputsize=300)

@app.post("/classify")
async def classifyimage(file: UploadFile = File(...)):

    """Classify an uploaded image."""
    contents = await file.read()
    image = np.array(Image.open(io.BytesIO(contents)).convert("RGB"))

    # Save temporarily and classify
    temppath = "/tmp/tempclassify.jpg"

    cv2.imwrite(temppath, cv2.cvtColor(image, cv2.COLORRGB2BGR))


    predictions = classifier.predict(temppath, top_k=5)
    return {"predictions": predictions}

Best Practices

Conclusion

Transfer learning is the most practical approach to image classification for the vast majority of real-world applications. By leveraging pre-trained models, you can achieve excellent accuracy with minimal data and compute. The workflow covered in this tutorial -- model selection, data preparation, fine-tuning with discriminative learning rates, rigorous evaluation, and ONNX export -- represents a production-ready pipeline applicable to any classification task.

Key takeaways: start simple (feature extraction), measure carefully (proper validation), and iterate (gradual unfreezing, hyperparameter tuning). With the timm library providing access to hundreds of state-of-the-art architectures, the bottleneck is no longer model availability but rather data quality and training methodology.