Image classification assigns a label (or a set of labels) to a whole image. It is the foundational computer vision task: most detection and segmentation backbones were first trained as classifiers, and transfer learning from ImageNet-pretrained classifiers is the default starting point for almost every vision project.
Parent topic: Computer Vision.
A CNN stacks convolution, normalisation, and nonlinearity layers in stages; each stage downsamples spatially and increases channels. Early layers respond to edges and colour blobs, later ones to object parts. Global average pooling collapses the final feature map to a vector, and a linear layer produces K logits. Softmax turns logits into probabilities: pk = ezk / Σj ezj, trained with cross-entropy L = −log py.
Common backbones: ResNet / EfficientNet, ConvNeXt, MobileNet (edge), and Vision Transformers. For multi-label problems (an image can be both "beach" and "sunset") replace softmax with per-class sigmoids and binary cross-entropy.
Training from scratch needs very large datasets. With a pretrained backbone, a few hundred labelled images per class is often enough for a useful model. Start with a linear probe as a baseline, then unfreeze progressively. Use a lower learning rate for pretrained layers than for the new head, and match the pretrained model's input normalisation.
import torch, torchvision
from torch import nn
from torchvision import datasets, transforms
from torchvision.models import resnet50, ResNet50_Weights
weights = ResNet50_Weights.IMAGENET1K_V2
train_tf = transforms.Compose([
transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(),
transforms.ColorJitter(0.2, 0.2, 0.2), transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
val_tf = weights.transforms()
# folder layout: data/train/<class>/*.jpg, data/val/<class>/*.jpg
train = datasets.ImageFolder("data/train", train_tf)
val = datasets.ImageFolder("data/val", val_tf)
tl = torch.utils.data.DataLoader(train, 32, shuffle=True, num_workers=4)
vl = torch.utils.data.DataLoader(val, 64, num_workers=4)
dev = "cuda" if torch.cuda.is_available() else "cpu"
model = resnet50(weights=weights)
for p in model.parameters():
p.requires_grad = False # linear probe first
model.fc = nn.Linear(model.fc.in_features, len(train.classes))
model.to(dev)
opt = torch.optim.AdamW(model.fc.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
for epoch in range(5):
model.train()
for x, y in tl:
x, y = x.to(dev), y.to(dev)
loss = loss_fn(model(x), y)
opt.zero_grad(); loss.backward(); opt.step()
model.eval(); correct = n = 0
with torch.no_grad():
for x, y in vl:
pred = model(x.to(dev)).argmax(1).cpu()
correct += (pred == y).sum().item(); n += len(y)
print(f"epoch {epoch}: val acc {correct / n:.3f}")After the probe converges, set requires_grad=True on layer4 (and later layer3) and continue with a learning rate around 10x lower.
Top-1 accuracy is fine for balanced benchmarks; for real data use per-class precision/recall, macro-F1, and the confusion matrix. Check calibration (do 0.9-confidence predictions succeed 90% of the time?) if downstream logic uses probabilities, and evaluate on data from the deployment environment, not only a random split of the training set.
Modern image classifiers are trained with a fairly standard recipe, and most gains on a real project come from applying it carefully rather than from a new architecture.
Data augmentation is the main regularizer. RandomResizedCrop (random scale and aspect ratio, then resize to 224) and horizontal flip are the baseline; color jitter, RandAugment or TrivialAugment add more variation. Mixup and CutMix blend images and labels, which improves robustness and calibration on many benchmarks. Choose augmentations that preserve the label: flipping is wrong for digits or left/right anatomy, and heavy color jitter is wrong when color is the signal (ripeness of fruit, skin lesions).
Optimization: AdamW or SGD with momentum, a short linear warmup and then cosine decay. For fine-tuning, a learning rate around 1e-4 to 1e-3 for the head and 10x lower for the backbone is a reasonable starting range; weight decay of 0.01–0.05 with AdamW is common. Label smoothing (target 0.9 for the true class, the remainder spread over others) reduces overconfidence. Mixed precision (bf16 or fp16 with loss scaling) roughly halves memory and speeds training on recent GPUs.
Input resolution is a real hyperparameter. Small objects or fine texture (defects, pathology tiles) need higher resolution or tiling; doubling the side length roughly quadruples compute for a CNN. Test-time augmentation (averaging predictions over a flip and a few crops) gives a modest, cheap accuracy gain when latency allows.
Suppose a 3-class model (cat, dog, fox) outputs logits z = (2.0, 1.0, 0.1) for an image of a cat. Exponentiating gives roughly (7.39, 2.72, 1.11), which sum to 11.22, so the softmax probabilities are about (0.66, 0.24, 0.10). The cross-entropy loss is −log(0.66) ≈ 0.42. If the image were actually a fox, the loss would be −log(0.10) ≈ 2.30: confident mistakes cost much more than uncertain ones, which is why the loss pushes probabilities toward honesty.
Deep networks are nonetheless often overconfident (Guo et al., 2017). Temperature scaling fixes much of this after training: divide every logit by a single scalar T fitted on a validation set. With T = 2 the same logits become (1.0, 0.5, 0.05), and the probabilities soften to about (0.48, 0.29, 0.23). The ranking, and therefore accuracy, does not change, only the confidence. Measure calibration with Expected Calibration Error (ECE) and a reliability diagram.
A softmax classifier must pick one of its K classes even for an image it has never seen anything like. In deployment that is a serious failure mode: a plant-disease model will confidently label a photo of a shoe. Simple baselines include thresholding the maximum softmax probability (Hendrycks & Gimpel, 2017), using the maximum logit or an energy score, and adding an explicit "other" class trained on diverse negative images. None is perfect, so route low-confidence predictions to a human or a safe default.
Other production concerns:
| Family | Strengths | Watch out for |
|---|---|---|
| ResNet-50 | Well understood, strong baseline, easy to fine-tune | Not the most accurate per FLOP today |
| EfficientNet / MobileNet | Good accuracy per parameter, mobile-friendly | Depthwise convs can be slower on some GPUs than FLOPs suggest |
| ConvNeXt | CNN modernized with transformer-era design choices | Larger variants need substantial compute |
| ViT / DeiT / Swin | Scale well with data and pretraining; strong transfer | Data-hungry from scratch; patch size sets resolution cost |
| CLIP image encoder | Zero-shot via text prompts; good linear probes | Inherits web-data biases; prompt wording matters |
A practical order: start with a linear probe on a CLIP or DINOv2 embedding (minutes, often a strong baseline), then fine-tune a ResNet-50 or ConvNeXt-T, and only then try larger models if the error analysis shows the model, not the data, is the bottleneck.