CNN Image Classification with PyTorch

CNN Image Classification with PyTorch solves a specific problem: conventional neural networks treat every pixel independently, destroying the spatial relationships that make images meaningful. CNNs preserve spatial context through convolution operations, making them the industry standard for computer vision.

The practical consequence: a 224x224 RGB image has 150,528 pixel values. A fully connected layer treating each pixel independently would need millions of parameters for the first layer alone — and those parameters would encode no understanding of which pixels are neighbors. CNNs solve this with weight sharing. The same kernel slides across the entire image, learning features that are useful regardless of where in the frame they appear. A vertical edge detector works in the top-left corner and the bottom-right corner. You train it once, not once per position.

The most common production failure I keep seeing, even in teams with strong ML backgrounds: developers build a CNN, hit solid training accuracy, ship it, and then watch it fall apart on real traffic. The model learned the training distribution. It learned that cats appear on white backgrounds in their dataset. It learned that defective parts are photographed under a specific lighting rig. Data augmentation is not a nice-to-have — it is the difference between a model that generalizes and a model that memorizes.

What Is CNN Image Classification with PyTorch and Why Does It Exist?

CNN Image Classification with PyTorch is PyTorch's answer to a very specific problem: the curse of dimensionality in image data. A standard fully connected network processing a 224x224 RGB image would need over 150,000 weights connecting to just the first neuron — and those weights would carry no structural knowledge that adjacent pixels are related. Rotate the image by one pixel and the entire activation pattern changes. The network has to relearn every spatial variant of every feature independently.

CNNs break that problem by introducing two architectural constraints: local receptive fields and weight sharing. A 3x3 convolutional kernel connects each output neuron to only a 3x3 patch of input pixels, not the full image. And the same kernel — the same nine weights — is reused at every spatial position. A filter that detects a vertical edge in the top-left corner detects it identically in the bottom-right corner, using zero additional parameters. This is what makes CNNs tractable for images.

The hierarchical structure is the other half of the story. Early layers detect the simplest visual primitives — edges, gradients, color transitions. Middle layers combine those primitives into shapes and textures. Deeper layers combine shapes into object parts. The final layers combine parts into full object representations. This mirrors how biological visual cortex is organized, though the analogy should not be over-read — CNNs converge on this structure because it works, not because they were explicitly designed to replicate neuroscience.

The performance trade-off worth understanding before you write a single line: CNNs are translation-invariant by design (the same feature is detected regardless of position) but they are not rotation-invariant or scale-invariant out of the box. A model trained only on upright dogs will struggle with dogs photographed at a 45-degree angle. Data augmentation — random rotations, flips, scaling — is what closes that gap. It is not optional and it is not a regularization afterthought. It is a core part of making the spatial inductive bias of CNNs work in practice.

import torch
import torch.nn as nn
import torch.nn.functional as F

# io.thecodeforge: Standard CNN architecture for image classification
# Input assumption: 3-channel RGB images at 224x224 resolution, 10-class output
class ForgeCNN(nn.Module):
    def __init__(self, num_classes: int = 10):
        super(ForgeCNN, self).__init__()

        # Block 1: 3 input channels -> 16 feature maps
        # padding=1 with kernel_size=3 preserves spatial dimensions (same padding)
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(16)  # stabilizes training, especially early on

        # Block 2: 16 feature maps -> 32 feature maps
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(32)

        # MaxPool2d with kernel=2, stride=2 halves both spatial dimensions
        # After pool1: 224 -> 112. After pool2: 112 -> 56.
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        # AdaptiveAvgPool2d fixes output to 7x7 regardless of input resolution
        # This eliminates the dimension recalculation problem when you change upstream layers
        self.adaptive_pool = nn.AdaptiveAvgPool2d(output_size=(7, 7))

        # Flattened size: 32 channels * 7 * 7 = 1568 — fixed regardless of input resolution
        self.fc1 = nn.Linear(32 * 7 * 7, 256)
        self.dropout = nn.Dropout(p=0.4)
        self.fc2 = nn.Linear(256, num_classes)

        # Verify dimensions at init time — fails loudly in development, not silently in production
        self._verify_dimensions()

    def _verify_dimensions(self):
        """Pass a dummy tensor at init to catch shape errors immediately."""
        with torch.no_grad():
            dummy = torch.randn(1, 3, 224, 224)
            out = self.forward(dummy)
            assert out.shape == (1, self.fc2.out_features), (
                f"Output shape mismatch: expected (1, {self.fc2.out_features}), got {out.shape}"
            )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Conv -> BN -> ReLU -> Pool
        x = self.pool(F.relu(self.bn1(self.conv1(x))))   # [B, 16, 112, 112]
        x = self.pool(F.relu(self.bn2(self.conv2(x))))   # [B, 32, 56, 56]
        x = self.adaptive_pool(x)                         # [B, 32, 7, 7] — fixed output
        x = x.view(x.size(0), -1)                         # [B, 1568]
        x = self.dropout(F.relu(self.fc1(x)))             # [B, 256]
        x = self.fc2(x)                                   # [B, num_classes]
        return x


# Instantiation — dimension verification runs automatically
model = ForgeCNN(num_classes=10)
print(f"Total parameters: {sum(p.numel() for p in model.parameters()):,}")

Enterprise Deployment: Scaling Vision Models with Docker

In a professional environment, deploying a CNN is not a question of zipping up a .pt file and calling it done. You need a reproducible runtime that guarantees the model running inference in production is operating under exactly the same conditions as the one you validated on your workstation — same PyTorch version, same CUDA version, same cuDNN version. The gap between those environments is where silent failures live.

The containerization strategy for vision services has become fairly standardized: use the official PyTorch Docker images as your base, pin the exact version triple (PyTorch, CUDA, cuDNN), copy only the inference code and model weights into the image, and run as a non-root user. Nothing else belongs in a production image.

The failure mode I see most often in teams that have not gone through this before: the model trains on a workstation with CUDA 12.1 and cuDNN 8.9. The production cluster was provisioned six months earlier and runs CUDA 11.8 with cuDNN 8.6. The model loads, inference completes, but the predictions are subtly wrong — not crashed, not obviously broken, just quietly wrong. cuDNN version differences can produce numerically different outputs for the same input through the same weights. The divergence is small enough that unit tests pass but large enough to affect classification confidence scores. Pinning versions in the Dockerfile is not pedantry — it is the only way to make this class of failure impossible.

For teams running multiple vision services, the image size difference between runtime and devel images matters at scale. A devel image carrying a full CUDA compiler toolchain is typically 6-8GB. A runtime image is 2-3GB. When you are pulling that image across dozens of nodes during a rolling deployment, the difference is not trivial.

# io.thecodeforge: Production Vision Service Environment
# Pin the full version triple: PyTorch + CUDA + cuDNN
# Never use 'latest' — it will silently change and break reproducibility
FROM pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime

WORKDIR /app

# Install Python dependencies first — this layer caches independently of source code
# Avoids re-installing packages on every source change
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy model weights and inference source separately
# Model weights change less frequently than source code
# Keeping them in separate COPY instructions preserves Docker layer cache
COPY ./models /app/models
COPY ./src /app/src

# Run as a non-root user — required by most enterprise security policies
# and good practice regardless
USER 1000

# Health check: verify the inference service starts and the model loads correctly
HEALTHCHECK --interval=30s --timeout=10s --retries=3 \
    CMD curl -f http://localhost:8080/health || exit 1

ENTRYPOINT ["python", "src/inference_service.py"]

Data Persistence: Tracking Image Metadata in SQL

Vision projects at any serious scale involve millions of images. Storing binary image data in a relational database is almost always the wrong call — databases are optimized for structured queries, not serving multi-megabyte blobs. The correct architecture is a clean separation: SQL tracks file paths, labels, split assignments, and metadata. The actual image files live on disk or in object storage like S3 or GCS. The PyTorch DataLoader queries SQL for the current split, gets back a list of paths, and loads images from disk on demand with transforms applied per batch.

This pattern solves several real problems that teams hit as their datasets grow. When you add new images, you insert a database row and drop the file in the right location — the next training run picks them up automatically, correctly assigned to their split. Without this, teams maintain CSV files that drift from the actual filesystem over time. Someone adds images, forgets to update the CSV, and the next training run silently ignores a quarter of the new data.

The split assignment in SQL also gives you something CSV files cannot easily give you: reproducible experiments. You can query for exactly which images were in the validation set for experiment run 47, long after the fact. You can audit class balance per split. You can add a 'held_out' split for images you want to exclude from training without deleting them. The metadata layer is cheap to maintain and it pays back continuously throughout the lifetime of a project.

One thing to get right from the start: store file paths as relative paths or as object storage keys, not as absolute filesystem paths. Absolute paths break the moment you move the dataset to a different machine, mount it at a different path, or migrate from local disk to S3.

-- io.thecodeforge: Schema for managing training image references
-- Store paths and labels here. Store actual image bytes on disk or S3.
CREATE TABLE IF NOT EXISTS io.thecodeforge.vision_assets (
    asset_id          UUID        PRIMARY KEY DEFAULT gen_random_uuid(),
    file_path         TEXT        NOT NULL UNIQUE,  -- relative path or S3 key, never absolute
    label_id          INT         NOT NULL,
    label_name        VARCHAR(128),                 -- denormalized for readability in queries
    split_type        VARCHAR(10) CHECK (split_type IN ('train', 'val', 'test', 'held_out')),
    source_dataset    VARCHAR(64),                  -- track provenance when merging multiple datasets
    is_augmented      BOOLEAN     DEFAULT FALSE,    -- flag synthetically generated images
    last_augmented_at TIMESTAMP,
    created_at        TIMESTAMP   DEFAULT NOW()
);

-- Index on split_type for fast DataLoader queries
CREATE INDEX IF NOT EXISTS idx_vision_assets_split
    ON io.thecodeforge.vision_assets (split_type, label_id);

-- Fetch a random training batch — used by the PyTorch DataLoader during epoch iteration
SELECT file_path, label_id
FROM   io.thecodeforge.vision_assets
WHERE  split_type = 'train'
  AND  is_augmented = FALSE           -- exclude pre-generated augmented copies if present
ORDER  BY RANDOM()
LIMIT  32;

-- Audit class balance before training — imbalanced classes need weighted sampling
SELECT label_name, split_type, COUNT(*) AS image_count
FROM   io.thecodeforge.vision_assets
GROUP  BY label_name, split_type
ORDER  BY label_name, split_type;

Common Mistakes and How to Avoid Them

Most CNN bugs are not subtle. They fall into a small set of categories that you see over and over once you have reviewed enough ML codebases. The frustrating part is that many of them do not produce errors — they produce a model that trains, passes validation, and then quietly fails in production.

The dimension mismatch between the last conv layer and the first linear layer is the most common hard error. PyTorch constructs the computational graph dynamically, which means it does not validate the connection between your conv stack and your linear layer when you define the model — only when you run a forward pass. If you never run a forward pass on the full model during development (easy to miss if you are training in a notebook and checking only individual layer outputs), the mismatch survives until inference.

A subtler mistake: calling model.forward(x) directly instead of model(x). The difference is that model(x) goes through the nn.Module __call__ mechanism, which fires all registered forward hooks. Profilers, debuggers, gradient checkpointing, and libraries like torchvision's feature extraction API all rely on these hooks. Calling forward() directly bypasses them. The output is numerically identical but you silently opt out of the entire hook infrastructure.

The production mistake that costs the most: developers normalize training images with ImageNet statistics (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) during training, but the normalization transform is defined inline in the training script rather than stored alongside the model weights. Six months later, someone writes a new inference service, does not realize normalization is required, and ships it. The model receives raw pixel values in the range [0, 255] or [0.0, 1.0] when it was trained on normalized inputs. Predictions are garbage. The fix is to store normalization parameters in the model checkpoint and load them in every inference pipeline — treat them as part of the model artifact, not as a training detail.

# io.thecodeforge: Common mistake patterns and their correct alternatives
import torch
import torch.nn as nn
from torchvision import transforms, models


# ─── MISTAKE 1: Hardcoded flattened dimension ────────────────────────────────
# Bad: fails silently until forward pass if you change upstream layers
class BadCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        # This hardcoded number is wrong if you add another pooling layer
        self.classifier = nn.Linear(32 * 112 * 112, 10)


# Good: use AdaptiveAvgPool2d — flattened size is always channels * H * W
class GoodCNN(nn.Module):
    def __init__(self, num_classes: int = 10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.adaptive_pool = nn.AdaptiveAvgPool2d((7, 7))  # fixed output: 32 * 7 * 7 = 1568
        self.classifier = nn.Linear(32 * 7 * 7, num_classes)

    def forward(self, x):
        x = self.features(x)
        x = self.adaptive_pool(x)
        return self.classifier(x.view(x.size(0), -1))


# ─── MISTAKE 2: Forgetting to store normalization with the model ─────────────
# Bad: normalization parameters live only in the training script
# train_transforms = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
# ...six months later, the inference service ships without normalization

# Good: save normalization params in the checkpoint
NORM_MEAN = [0.485, 0.456, 0.406]
NORM_STD = [0.229, 0.224, 0.225]

def save_model_with_norm(model, path):
    torch.save({
        'state_dict': model.state_dict(),
        'norm_mean': NORM_MEAN,
        'norm_std': NORM_STD,
    }, path)

def load_model_with_norm(model_class, path):
    checkpoint = torch.load(path, map_location='cpu')
    model = model_class()
    model.load_state_dict(checkpoint['state_dict'])
    inference_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=checkpoint['norm_mean'], std=checkpoint['norm_std']),
    ])
    return model, inference_transform


# ─── MISTAKE 3: Calling model.forward() directly ────────────────────────────
# Bad: bypasses all registered forward hooks
# output = model.forward(input_tensor)

# Good: call the model as a callable
# output = model(input_tensor)


# ─── MISTAKE 4: Using augmentation during validation ────────────────────────
# Bad: validation results are non-deterministic, benchmark comparisons are meaningless
bad_val_transforms = transforms.Compose([
    transforms.RandomHorizontalFlip(),  # should never be in a validation pipeline
    transforms.ToTensor(),
])

# Good: validation pipeline is fully deterministic — no random transforms
good_val_transforms = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=NORM_MEAN, std=NORM_STD),
])

Transfer Learning: The Production Default for CNNs

Transfer learning is the single most impactful technique in practical computer vision, and it is underused in proportion to how well it is understood. The idea is simple: a model pre-trained on ImageNet has already learned to detect edges, textures, shapes, and object parts from 1.2 million labeled images. Those features are not specific to ImageNet — they are general visual features that appear in almost every real-world image domain. Instead of spending compute and data teaching a new model what an edge is, you start from weights that already know, and you teach only the final classification decision.

The practical numbers matter here. A ResNet-18 pre-trained on ImageNet and fine-tuned on a new 1,000-image-per-class dataset will typically reach 88-92% validation accuracy in 10-15 epochs. The same architecture trained from random initialization on the same 1,000 images per class will take hundreds of epochs to converge and will likely overfit to 70-75% accuracy despite all regularization efforts. The pre-trained model is not just faster to train — it learns better because the early layers are already at a strong initialization point.

The fine-tuning strategy has two stages and both matter. In the warm-up stage, you freeze all layers except the final classification head and train for 5-10 epochs. This lets the new classifier head converge without large gradients propagating through the pre-trained features and destroying them — a phenomenon sometimes called feature forgetting. In the fine-tuning stage, you unfreeze all layers and continue training with a learning rate that is roughly 10x lower than the warm-up rate. The lower rate allows the pre-trained features to adapt to your domain without being overwritten.

The one case where transfer learning from ImageNet genuinely does not work well: domains where the low-level visual statistics are fundamentally different from natural images. Medical X-rays, radar imagery, spectrograms, and electron microscopy images have different edge distributions, texture statistics, and spatial structures than photographs of objects. In these cases, you often get better results fine-tuning with all layers unfrozen from the start, or training from scratch if you have enough domain data.

# io.thecodeforge: Production transfer learning pattern with two-stage fine-tuning
import torch
import torch.nn as nn
from torchvision import models

NUM_CLASSES = 10
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


# ─── Stage 0: Load pre-trained weights ──────────────────────────────────────
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)

# ResNet-18 architecture breakdown:
# - 17 convolutional layers learning features from edges up to object parts
# - Final fc layer: Linear(512, 1000) for ImageNet 1000 classes
print(f"Original output classes: {model.fc.out_features}")


# ─── Stage 1: Freeze backbone, train head only (warm-up) ────────────────────
for param in model.parameters():
    param.requires_grad = False

# Replace the classifier head — this new layer has requires_grad=True by default
model.fc = nn.Sequential(
    nn.Linear(512, 256),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(256, NUM_CLASSES),
)

model = model.to(DEVICE)

# Only the new head parameters are optimized during warm-up
warmup_optimizer = torch.optim.Adam(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=1e-3
)

print(f"Warm-up trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
print(f"Total params: {sum(p.numel() for p in model.parameters()):,}")

# Train for 5-10 warm-up epochs here (training loop omitted for brevity)
# ...after warm-up epochs...


# ─── Stage 2: Unfreeze all layers and fine-tune with lower LR ───────────────
for param in model.parameters():
    param.requires_grad = True

# Lower learning rate prevents large gradients from destroying pre-trained features
# Layer-wise learning rate decay: backbone gets 10x smaller LR than the head
backbone_params = [p for n, p in model.named_parameters() if 'fc' not in n]
head_params     = [p for n, p in model.named_parameters() if 'fc' in n]

finetune_optimizer = torch.optim.Adam([
    {'params': backbone_params, 'lr': 1e-5},   # backbone: small LR, preserve features
    {'params': head_params,     'lr': 1e-4},   # head: larger LR, learn class boundaries
])

# Cosine annealing reduces LR smoothly over fine-tuning epochs
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    finetune_optimizer, T_max=20, eta_min=1e-7
)

print(f"Fine-tuning trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")

Data Augmentation: The Regularization You Cannot Skip

Data augmentation is the most consistently underestimated technique in CNN training. Developers new to computer vision tend to treat it as a preprocessing detail — something to add if you have time, after you get the architecture working. That is exactly backwards. Augmentation is the primary mechanism by which CNNs learn invariance to transformations that do not change object identity, and without it, your model is almost guaranteed to memorize training-specific artifacts rather than learn generalizable features.

The core insight: every time augmentation applies a random horizontal flip to an image, the model sees what is essentially a new training example. The label does not change — a cat is still a cat whether it faces left or right — but the pixel values are different. Over thousands of batches, this teaches the model that left-right orientation is not a distinguishing feature of the object class. The same logic applies to every other augmentation: random crops teach position invariance, color jitter teaches lighting invariance, random erasing prevents the model from relying on a single dominant texture patch.

The augmentation strategy is domain-dependent and this is where teams make mistakes. Horizontal flip is safe for natural images of objects, vehicles, animals, and most industrial inspection tasks. It is not safe for medical images where laterality matters — a left lung X-ray is not equivalent to a right lung X-ray. It is not safe for text-in-image classification tasks where mirrored text is not valid text. Random rotation up to 15-20 degrees is generally safe. Rotation up to 180 degrees is not safe unless your objects genuinely appear at all orientations in production (aerial imagery being a common exception).

The pipeline separation is non-negotiable: training gets augmentation, validation and inference get only deterministic transforms. Applying augmentation during validation makes your benchmark numbers non-reproducible. Applying augmentation during inference makes your predictions non-deterministic — the same image sent twice gets different predictions. Both are unacceptable in a production system.

# io.thecodeforge: Domain-aware augmentation pipelines for 2026 production use
from torchvision import transforms
import torchvision.transforms.v2 as transforms_v2  # v2 API: faster, supports bounding boxes

NORM_MEAN = [0.485, 0.456, 0.406]
NORM_STD  = [0.229, 0.224, 0.225]


# ─── Natural images (objects, products, vehicles) ────────────────────────────
train_transforms = transforms_v2.Compose([
    transforms_v2.RandomResizedCrop(224, scale=(0.7, 1.0), ratio=(0.75, 1.33)),
    transforms_v2.RandomHorizontalFlip(p=0.5),
    transforms_v2.RandomRotation(degrees=15),
    transforms_v2.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.05),
    transforms_v2.RandomGrayscale(p=0.05),          # occasional grayscale improves robustness
    transforms_v2.ToImage(),
    transforms_v2.ToDtype(torch.float32, scale=True),
    transforms_v2.Normalize(mean=NORM_MEAN, std=NORM_STD),
    transforms_v2.RandomErasing(p=0.1, scale=(0.02, 0.15)),  # cutout regularization
])


# ─── Validation and inference — always deterministic ─────────────────────────
val_transforms = transforms_v2.Compose([
    transforms_v2.Resize(256),
    transforms_v2.CenterCrop(224),
    transforms_v2.ToImage(),
    transforms_v2.ToDtype(torch.float32, scale=True),
    transforms_v2.Normalize(mean=NORM_MEAN, std=NORM_STD),
])
# Important: val_transforms == inference_transforms
# If they differ, your benchmark does not represent production behavior


# ─── Medical imaging — laterality-aware, no horizontal flip ──────────────────
medical_train_transforms = transforms_v2.Compose([
    transforms_v2.RandomResizedCrop(224, scale=(0.85, 1.0)),
    transforms_v2.RandomRotation(degrees=10),       # small rotation only
    transforms_v2.ColorJitter(brightness=0.15, contrast=0.15),  # conservative color shift
    # NO RandomHorizontalFlip — left/right matters for anatomical laterality
    transforms_v2.ToImage(),
    transforms_v2.ToDtype(torch.float32, scale=True),
    transforms_v2.Normalize(mean=NORM_MEAN, std=NORM_STD),
])


# ─── Satellite / aerial imagery — rotation-invariant by nature ───────────────
aerial_train_transforms = transforms_v2.Compose([
    transforms_v2.RandomResizedCrop(224, scale=(0.7, 1.0)),
    transforms_v2.RandomHorizontalFlip(p=0.5),
    transforms_v2.RandomVerticalFlip(p=0.5),        # valid for top-down imagery
    transforms_v2.RandomRotation(degrees=90),       # full rotation invariance for aerial views
    transforms_v2.ColorJitter(brightness=0.2, contrast=0.2),
    transforms_v2.ToImage(),
    transforms_v2.ToDtype(torch.float32, scale=True),
    transforms_v2.Normalize(mean=NORM_MEAN, std=NORM_STD),
])

Frequently Asked Questions