Image Classification with Deep Learning

Module: Computer Vision | Difficulty: Advanced

Classification as Probabilistic Inference

Image classification maps an input image to one of predefined categories. Modern deep learning approaches treat this as probabilistic inference, where a neural network outputs a probability distribution over classes. The network learns a mapping from image space to the -simplex of probability distributions.

The challenge lies in learning representations that are invariant to irrelevant variations (lighting, pose, background) while being discriminative for the target classes. Deep CNNs address this through hierarchical feature extraction, where early layers detect universal visual patterns (edges, textures) and deeper layers compose these into class-specific features.

Softmax Output

The softmax function converts raw logits to a valid probability distribution:

Where each parameter means:

• — raw output (logit) for class
• — predicted probability for class
• — total number of classes
• Intuition: Softmax exponentiates each logit to make it positive, then normalizes by the sum to ensure all probabilities sum to 1; larger logits produce higher probabilities

Cross-Entropy Loss

The standard training objective minimizes the negative log-likelihood:

Where each parameter means:

• — one-hot encoded ground truth (1 for correct class, 0 otherwise)
• — predicted probability for class
• Intuition: For the correct class, this simplifies to , which approaches 0 as confidence increases and approaches infinity as confidence decreases

Label Smoothing

Label smoothing regularizes the model by softening target distributions:

Where each parameter means:

• — smoothing parameter (typically 0.1)
• — original one-hot label
• — smoothed label distribution
• Intuition: Instead of requiring the model to be 100% confident, label smoothing encourages it to be 90% confident, preventing overconfident predictions and improving calibration

Data Augmentation

Data augmentation artificially expands the training set by applying transformations:

Where each parameter means:

• — random transformation (rotation, flip, crop, color jitter)
• — augmented image
• — same label (transformation-invariant)
• Intuition: By showing the model many variations of each image, it learns to ignore irrelevant differences and focus on discriminative features

Residual Learning

Skip Connections

Deep networks suffer from degradation: adding more layers increases training error. ResNet addresses this by learning residuals instead of direct mappings:

Where each parameter means:

• — desired mapping
• — residual function to learn
• — identity skip connection
• Intuition: If the optimal mapping is close to identity, it's easier for the network to learn than to learn directly; skip connections provide a gradient highway that prevents vanishing gradients

Compound Scaling

EfficientNet scales depth, width, and resolution together using a compound coefficient:

subject to .

Where each parameter means:

• — compound scaling coefficient (user-specified)
• — coefficients determined by grid search
• Intuition: Scaling all three dimensions together maintains balance; adding depth without resolution wastes compute, while adding resolution without depth limits the receptive field

Classification Model Comparison

Complete ResNet Fine-Tuning Pipeline

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


def create_finetuned_model(num_classes, pretrained=True):
    model = models.resnet50(pretrained=pretrained)
    for param in model.parameters():
        param.requires_grad = False
    num_features = model.fc.in_features
    model.fc = nn.Sequential(
        nn.Dropout(0.3),
        nn.Linear(num_features, 512),
        nn.ReLU(inplace=True),
        nn.Dropout(0.2),
        nn.Linear(512, num_classes)
    )
    return model


def get_transforms(image_size=224):
    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(image_size),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    val_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(image_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    return train_transform, val_transform


def train_model(model, train_loader, val_loader, epochs=20, lr=0.001):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
    optimizer = optim.AdamW(model.fc.parameters(), lr=lr, weight_decay=0.01)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    best_acc = 0.0
    for epoch in range(epochs):
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
        train_acc = 100.0 * correct / total
        model.eval()
        val_correct = 0
        val_total = 0
        with torch.no_grad():
            for images, labels in val_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                _, predicted = outputs.max(1)
                val_total += labels.size(0)
                val_correct += predicted.eq(labels).sum().item()
        val_acc = 100.0 * val_correct / val_total
        scheduler.step()
        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), 'best_model.pth')
        print(f"Epoch {epoch+1}/{epochs}: "
              f"Train Acc: {train_acc:.1f}%, Val Acc: {val_acc:.1f}%")
    return model


train_transform, val_transform = get_transforms()
model = create_finetuned_model(num_classes=10, pretrained=True)
params = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Total: {params:,}, Trainable: {trainable:,}")

Common Challenges

1. Overfitting on Small Datasets: Limited training data causes models to memorize rather than generalize, requiring strong augmentation or regularization
2. Class Imbalance: Uneven class distributions bias models toward majority classes, requiring class-weighted losses or resampling
3. Domain Shift: Models trained on one distribution may fail on another (e.g., photos vs. sketches), requiring domain adaptation
4. Computational Cost: Training large models requires significant GPU resources and time, motivating efficient architectures
5. Calibration: Models often produce overconfident predictions, requiring temperature scaling or label smoothing

Case Study: Fine-Grained Classification

A 2020 study on bird species classification (CUB-200-2011, 11,788 images, 200 species) achieved 91.7% top-1 accuracy using EfficientNet-B4 with cutout and cutmix augmentation. The model was pretrained on ImageNet (1.28M images) and fine-tuned for 100 epochs with cosine annealing (initial LR=0.001). Key engineering decisions included: using larger input resolution (380×380), applying Test Time Augmentation (10 crops + flips), and using label smoothing (ε=0.1). The error analysis revealed that 62% of misclassifications involved visually similar species differing only in subtle plumage patterns, suggesting that attention mechanisms could further improve performance.

Advanced Data Augmentation Techniques

RandAugment

RandAugment automatically selects augmentation operations from a predefined set, reducing the need for extensive hyperparameter tuning. It applies N randomly chosen operations with magnitude M:

Where each parameter means:

• — number of randomly selected operations per image (typically 2-3)
• — magnitude of all operations (typically 9-15 on a scale of 0-30)
• — randomly selected transformation (rotation, shear, translate, contrast, etc.)
• Intuition: With N=2 and M=12, each image gets two random augmentations applied at moderate strength; this simple search space (2 hyperparameters) matches complex policies like AutoAugment

CutMix and MixUp

CutMix pastes rectangular patches from one image onto another, mixing labels proportionally:

Where each parameter means:

• — binary mask defining the cutout region
• — mixing ratio (uniform in [0, 1])
• — two randomly selected images
• — corresponding soft labels
• Intuition: CutMix encourages the model to recognize objects from partial views, improving localization ability while providing regularization

Random Erasing

Random erasing selects a rectangular region and replaces it with random values or mean pixel values:

Where each parameter means:

• — randomly selected rectangular region
• — replacement value (random or mean pixel value)
• Intuition: By randomly occluding parts of the image, the model learns to use multiple discriminative features rather than relying on a single region

Test-Time Augmentation (TTA)

TTA improves inference accuracy by applying augmentations to test images and averaging predictions:

Where each parameter means:

• — number of augmented versions (typically 5-10)
• — augmentation transformation (flip, multi-crop, color jitter)
• — model output probability for class
• Intuition: By averaging predictions over multiple views, TTA reduces variance and improves robustness to spatial variations; typically provides 0.5-1.5% accuracy boost

Model Ensembling

Ensemble methods combine multiple models to improve robustness:

Where each parameter means:

• — number of models in the ensemble
• — weight for model (typically equal weights)
• — output of model for class
• Intuition: Different models make different errors; ensembling averages out these errors, typically improving accuracy by 1-3% at the cost of M times more computation

Model Distillation

Knowledge distillation trains a smaller student model to mimic a larger teacher model:

Where each parameter means:

• — teacher logits divided by temperature
• — student logits divided by temperature
• — temperature parameter (typically 3-5)
• — balance weight (typically 0.1-0.5)
• Intuition: The teacher's soft labels provide richer information than hard labels, revealing class similarities; the student learns these relationships, achieving higher accuracy than training on hard labels alone

Key Takeaways

• Image classification maps images to probability distributions over K classes using softmax + cross-entropy
• Residual connections solve the degradation problem, enabling training of 100+ layer networks
• Transfer learning from ImageNet pretrained models is standard practice for most vision tasks
• Compound scaling of depth, width, and resolution improves efficiency beyond individual scaling
• Data augmentation is critical for regularization, with mixup and cutmix providing the largest gains
• Label smoothing improves calibration by preventing overconfident predictions
• Test-time augmentation and ensembling provide additional accuracy improvements at inference cost