Image Classification Tutorial: PyTorch Transfer Learning and Training
Understanding Image Classification
Image classification is the most fundamental computer vision task: assigning a single semantic label to an entire image. Despite its apparent simplicity, classification serves as the foundation for almost all other vision tasks. Object detection builds on classification by adding localization. Image segmentation extends classification to every pixel. Even complex tasks like visual question answering rely on classification models as backbone feature extractors.
The modern classification pipeline follows a standard pattern: data preparation, model selection (typically with transfer learning), training with augmentation and regularization, evaluation against held-out data, and deployment optimization. The ImageNet challenge, which drove classification progress from 2010 to 2017, saw error rates drop from 28% to 2.3% as architectures evolved from shallow networks to deep ResNets and finally Vision Transformers.
Data Preparation
Proper data preparation is the most impactful step. Images should be organized into a directory structure compatible with PyTorch’s ImageFolder dataset: separate train and val directories, each containing one subdirectory per class. The PyTorch documentation on data loading provides detailed guidance on this structure.
from torchvision import datasets, transforms
train_dir = 'data/train'
val_dir = 'data/val'
train_dataset = datasets.ImageFolder(
train_dir,
transform=transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])
)
val_dataset = datasets.ImageFolder(
val_dir,
transform=transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])
)The normalization values are the ImageNet dataset mean and standard deviation, which should be used when fine-tuning ImageNet-pretrained models. For custom datasets, compute dataset-specific statistics by iterating through the entire training set. Dataset imbalance — where some classes have many more examples than others — should be addressed through weighted sampling or class-balanced loss functions.
Transfer Learning
Transfer learning is the single most effective technique for training classification models with limited data. Instead of training from scratch (which requires hundreds of thousands of images), we take a model pre-trained on ImageNet and fine-tune it on our target dataset. The pre-trained model has already learned general visual features — edges, textures, shapes, object parts — that transfer to any image classification task.
The standard approach freezes the early layers (which capture generic features) and fine-tunes the later layers (which capture task-specific features). In PyTorch, this is implemented by setting requires_grad=False for all parameters, then re-initializing the final fully connected layer for the new number of classes:
import torchvision.models as models
model = models.resnet18(weights='IMAGENET1K_V1')
for param in model.parameters():
param.requires_grad = False
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2) # 2 classes for cat/dogFor larger datasets (1,000+ images per class), unfreeze all layers and train the entire network with a lower learning rate. For smaller datasets, freeze more layers — or use a linear classifier on top of frozen features. The general rule: the smaller the target dataset, the more aggressive the freezing.
Data Augmentation
Data augmentation generates additional training examples through label-preserving transformations. Each augmented image teaches the model to be invariant to that transformation. The standard augmentation pipeline for ImageNet-style classification includes random resized cropping (teaches scale and translation invariance), random horizontal flip (teaches mirror invariance), and color jittering (teaches lighting invariance).
train_transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ColorJitter(
brightness=0.2,
contrast=0.2,
saturation=0.2,
hue=0.1
),
transforms.RandomRotation(degrees=15),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])Advanced augmentation techniques push performance further. RandAugment randomly selects a subset of augmentation operations (translate, shear, rotate, solarize, equalize) with a learned magnitude parameter. MixUp creates weighted combinations of image pairs and their labels, forcing the model to learn linear interpolations between classes. CutMix cuts a patch from one image and pastes it onto another, combining region-level and image-level mixing.
Training Loop and Loss
The standard training loop iterates over epochs, each epoch processing the entire training dataset in batches. Cross-entropy loss is the standard for multi-class classification. The optimizer is typically SGD with momentum (0.9) and weight decay (1e-4), though AdamW has become increasingly popular for transformer-based architectures.
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
for inputs, labels in train_loader:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
scheduler.step()Learning Rate Scheduling
The learning rate schedule significantly impacts convergence quality. Cosine annealing reduces the learning rate following a cosine curve from the initial value to near zero. This schedule has no hyperparameters beyond the initial LR and produces the best average performance across architectures. OneCycleLR implements the cyclical learning rate policy: warm up from a low LR to a maximum, then cosine decay to near zero. The OneCycle policy often achieves higher accuracy in fewer epochs by exploring the loss landscape more thoroughly.
ReduceLROnPlateau reduces the learning rate by a factor when validation loss plateaus, providing adaptive scheduling that responds to actual training progress. The patience parameter controls how many epochs to wait before reducing the rate. Learning rate warmup — gradually increasing from a very small value to the target LR over the first 5-10 epochs — prevents early training instability, particularly with large batch sizes and transformer architectures.
Evaluation and Metrics
Classification evaluation uses several complementary metrics. Accuracy measures the fraction of correctly classified images but is misleading for imbalanced datasets. A model that always predicts the majority class achieves high accuracy while being useless. Per-class precision, recall, and F1-score provide a more nuanced view. The confusion matrix visualizes which classes are commonly confused.
Model Ensembling
Ensembling combines predictions from multiple independently trained models to improve accuracy and reduce variance. The simplest approach averages the softmax probabilities from 3-5 models trained with different random seeds. This typically improves accuracy by 1-3% over the best single model. Snapshot ensembling trains a single model with a cyclic learning rate, saving model checkpoints at each cycle’s minimum. The resulting ensemble captures diverse local minima without the cost of training multiple models.
Knowledge distillation can compress an ensemble into a single model. The student model is trained on the soft labels produced by the teacher ensemble, capturing the ensemble’s knowledge in a single forward pass at inference time. This technique is widely used in production to achieve ensemble-level accuracy with single-model inference speed.
Handling Edge Cases
Real-world deployment reveals edge cases that validation sets miss. Out-of-distribution detection identifies inputs that differ significantly from training data — the model should refuse to classify unknown categories rather than producing a confident wrong answer. Temperature scaling calibrates model confidence by dividing logits by a learned temperature parameter before softmax, ensuring that 90% confident predictions are actually correct 90% of the time.
Uncertainty estimation quantifies prediction reliability. Monte Carlo Dropout applies dropout at inference time, running multiple forward passes and measuring prediction variance. High variance indicates uncertainty. Deep Ensembles train multiple models and treat their prediction variance as uncertainty. These techniques are increasingly required for high-stakes applications in medical diagnosis and autonomous systems.
from sklearn.metrics import classification_report, confusion_matrix
model.eval()
all_preds = []
all_labels = []
with torch.no_grad():
for inputs, labels in val_loader:
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
all_preds.extend(preds.cpu().numpy())
all_labels.extend(labels.cpu().numpy())
print(classification_report(all_labels, all_preds))Model Deployment
After evaluation, the model is exported for deployment. The standard workflow converts the PyTorch model to TorchScript or ONNX format for production inference. TorchScript enables C++ inference without Python dependencies. ONNX provides cross-platform compatibility with TensorRT, OpenVINO, and ONNX Runtime. Quantization to FP16 or INT8 reduces model size by 2-4x with minimal accuracy loss, enabling deployment on edge devices.
# Export to ONNX
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, "model.onnx",
input_names=['input'],
output_names=['output'],
dynamic_axes={'input': {0: 'batch_size'},
'output': {0: 'batch_size'}})Frequently Asked Questions
How many images per class do I need? With transfer learning from ImageNet, 100-500 images per class can produce good results (85-95% accuracy depending on task difficulty). Training from scratch typically requires 1,000+ images per class. For very small datasets (10-50 images per class), consider using a frozen feature extractor with a linear classifier.
Which pretrained model should I use? ResNet-18 for speed and simplicity, ResNet-50 for balanced accuracy-efficiency, EfficientNet-B0 for best accuracy per parameter, and ViT-B/16 for maximum accuracy with sufficient data. For edge deployment, MobileNetV3 or EfficientNet-Lite offer the best size-accuracy trade-off.
Why is my training loss decreasing but validation accuracy not improving? This typically indicates overfitting — the model is memorizing training examples rather than learning generalizable features. Solutions: add more augmentation, increase dropout, reduce model capacity, or add weight decay.
What batch size should I use? Larger batch sizes (64-256) provide more stable gradients and faster training through parallelism. However, very large batches can lead to sharper minima that generalize worse. The learning rate should be scaled with batch size (linear scaling rule: double LR when doubling batch size). Use gradient accumulation to simulate large batches when GPU memory is limited.
How do I handle class imbalance? Weighted random sampling (WeightedRandomSampler in PyTorch) oversamples minority classes. Weighted loss functions (pass class weights to CrossEntropyLoss) penalize mistakes on minority classes more heavily. Focal loss is effective for extreme imbalance by down-weighting well-classified examples.
Related Articles
- Convolutional Neural Networks — architecture deep dive for classification models
- Object Detection Guide — extending classification to localization
- Computer Vision Datasets — data preparation and augmentation best practices