Computer vision with PyTorch has become the dominant approach for image classification, object detection, and visual recognition tasks in 2026. This tutorial covers building CNNs from scratch, applying transfer learning with pretrained models, and preparing for technical interview questions.

PyTorch 2.4 Performance

PyTorch 2.4 introduces torch.compile() by default for CNN operations, delivering 30-50% speedups on modern GPUs without code changes. All examples in this article are compatible with PyTorch 2.4+.

Understanding Convolutional Neural Networks for Image Classification

Convolutional Neural Networks extract hierarchical features from images through learnable filters. Unlike fully connected networks, CNNs preserve spatial relationships between pixels, making them ideal for visual tasks. The architecture consists of convolutional layers that detect edges and patterns, pooling layers that reduce dimensionality, and fully connected layers that perform classification.

A convolution operation slides a small filter (typically 3x3 or 5x5) across the input image, computing dot products at each position. This produces feature maps highlighting specific patterns like edges, textures, or shapes. Deeper layers combine these low-level features into complex representations—faces, objects, or scenes.

python

# cnn_architecture.py
import torch
import torch.nn as nn

class SimpleCNN(nn.Module):
    """Basic CNN for CIFAR-10 classification (32x32 RGB images, 10 classes)"""
    
    def __init__(self, num_classes: int = 10):
        super().__init__()
        # First conv block: 3 input channels (RGB) -> 32 feature maps
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),  # Output: 32x32x32
            nn.BatchNorm2d(32),  # Normalize activations for stable training
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2)  # Output: 16x16x32
        )
        # Second conv block: increase depth for richer features
        self.conv2 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=3, padding=1),  # Output: 16x16x64
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2)  # Output: 8x8x64
        )
        # Third conv block: capture high-level patterns
        self.conv3 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, padding=1),  # Output: 8x8x128
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2)  # Output: 4x4x128
        )
        # Classifier head
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 4 * 4, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),  # Prevent overfitting
            nn.Linear(256, num_classes)
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        return self.classifier(x)

This architecture progressively increases the number of filters (32 → 64 → 128) while reducing spatial dimensions through pooling. BatchNorm stabilizes training by normalizing layer outputs, while Dropout prevents the classifier from memorizing training samples.

Training a CNN with PyTorch DataLoaders

Efficient data loading is critical for GPU utilization during training. PyTorch's DataLoader handles batching, shuffling, and parallel data loading automatically. Proper data augmentation—random crops, flips, and color jitter—significantly improves generalization on unseen images.

python

# train_cnn.py
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Data augmentation for training (reduces overfitting)
train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomCrop(32, padding=4),  # Random crop with padding
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465),  # CIFAR-10 mean
                         (0.2470, 0.2435, 0.2616))  # CIFAR-10 std
])

# No augmentation for validation (deterministic evaluation)
val_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465),
                         (0.2470, 0.2435, 0.2616))
])

def train_model(model: nn.Module, epochs: int = 20) -> dict:
    """Train CNN with standard best practices"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    
    # Load CIFAR-10 dataset
    train_data = datasets.CIFAR10(root="./data", train=True, 
                                   download=True, transform=train_transform)
    val_data = datasets.CIFAR10(root="./data", train=False,
                                 transform=val_transform)
    
    # DataLoaders with num_workers for parallel loading
    train_loader = DataLoader(train_data, batch_size=128, shuffle=True,
                              num_workers=4, pin_memory=True)
    val_loader = DataLoader(val_data, batch_size=256, shuffle=False,
                            num_workers=4, pin_memory=True)
    
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    
    best_acc = 0.0
    for epoch in range(epochs):
        model.train()
        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()
        
        scheduler.step()
        
        # Validation
        model.eval()
        correct, total = 0, 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)
                total += labels.size(0)
                correct += predicted.eq(labels).sum().item()
        
        acc = 100. * correct / total
        best_acc = max(best_acc, acc)
        print(f"Epoch {epoch+1}/{epochs} - Val Acc: {acc:.2f}%")
    
    return {"best_accuracy": best_acc}

Key training optimizations include: pin_memory=True for faster CPU-to-GPU transfers, num_workers for parallel data loading, AdamW optimizer with weight decay for regularization, and cosine annealing learning rate schedule for smooth convergence.

Transfer Learning with Pretrained ResNet and EfficientNet

Transfer learning leverages models pretrained on ImageNet (1.2 million images, 1000 classes) and adapts them to custom datasets. This approach achieves higher accuracy with less training data and compute time. PyTorch's torchvision.models provides state-of-the-art architectures with pretrained weights.

The standard approach freezes early convolutional layers (which detect universal features like edges) and fine-tunes later layers plus a new classification head for the target task.

python

# transfer_learning.py
import torch
import torch.nn as nn
from torchvision import models
from torchvision.models import ResNet50_Weights, EfficientNet_B0_Weights

def create_transfer_model(
    model_name: str = "resnet50",
    num_classes: int = 10,
    freeze_backbone: bool = True
) -> nn.Module:
    """Create a pretrained model with custom classification head"""
    
    if model_name == "resnet50":
        # Load ResNet50 with ImageNet weights
        model = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
        # Replace final fully connected layer
        in_features = model.fc.in_features  # 2048 for ResNet50
        model.fc = nn.Sequential(
            nn.Dropout(0.3),
            nn.Linear(in_features, num_classes)
        )
        
    elif model_name == "efficientnet_b0":
        # EfficientNet: better accuracy/compute tradeoff
        model = models.efficientnet_b0(weights=EfficientNet_B0_Weights.IMAGENET1K_V1)
        in_features = model.classifier[1].in_features  # 1280 for B0
        model.classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(in_features, num_classes)
        )
    
    if freeze_backbone:
        # Freeze all layers except classifier
        for name, param in model.named_parameters():
            if "fc" not in name and "classifier" not in name:
                param.requires_grad = False
    
    return model

def count_trainable_params(model: nn.Module) -> int:
    """Count parameters that will be updated during training"""
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

# Example usage
model = create_transfer_model("resnet50", num_classes=5, freeze_backbone=True)
print(f"Trainable parameters: {count_trainable_params(model):,}")
# Output: ~10,245 (only classifier) vs ~25 million (full ResNet50)

Freezing the backbone reduces trainable parameters from 25 million to around 10,000, enabling training on small datasets without overfitting. For larger datasets (10,000+ images), gradually unfreeze later layers using differential learning rates—lower rates for pretrained layers, higher rates for the new classifier.

Ready to ace your Data Science & ML interviews?

Practice with our interactive simulators, flashcards, and technical tests.

Explore Data Science & ML

Data Augmentation Strategies for Computer Vision

Data augmentation artificially expands the training set by applying transformations that preserve semantic content. Modern augmentation goes beyond basic flips and rotations—techniques like MixUp and CutOut create synthetic training examples that improve model robustness.

python

# augmentation_strategies.py
import torch
import torchvision.transforms.v2 as T
from torchvision.transforms.v2 import functional as F

# Modern augmentation pipeline using torchvision v2 transforms
advanced_augmentation = T.Compose([
    T.RandomResizedCrop(224, scale=(0.8, 1.0)),  # Random crop and resize
    T.RandomHorizontalFlip(p=0.5),
    T.RandomVerticalFlip(p=0.1),  # Less common but useful for some domains
    T.RandomRotation(degrees=15),
    T.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1),
    T.RandomAffine(degrees=0, translate=(0.1, 0.1)),  # Small translations
    T.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),
    T.ToImage(),
    T.ToDtype(torch.float32, scale=True),
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

class CutOut:
    """Randomly mask out square regions of the image"""
    def __init__(self, n_holes: int = 1, length: int = 16):
        self.n_holes = n_holes
        self.length = length
    
    def __call__(self, img: torch.Tensor) -> torch.Tensor:
        h, w = img.shape[1], img.shape[2]
        mask = torch.ones_like(img)
        
        for _ in range(self.n_holes):
            y = torch.randint(0, h, (1,)).item()
            x = torch.randint(0, w, (1,)).item()
            y1 = max(0, y - self.length // 2)
            y2 = min(h, y + self.length // 2)
            x1 = max(0, x - self.length // 2)
            x2 = min(w, x + self.length // 2)
            mask[:, y1:y2, x1:x2] = 0
        
        return img * mask

def mixup_data(
    x: torch.Tensor,
    y: torch.Tensor,
    alpha: float = 0.4
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]:
    """MixUp: blend two images and their labels"""
    lam = torch.distributions.Beta(alpha, alpha).sample().item()
    batch_size = x.size(0)
    index = torch.randperm(batch_size)
    
    mixed_x = lam * x + (1 - lam) * x[index]
    y_a, y_b = y, y[index]
    
    return mixed_x, y_a, y_b, lam

CutOut forces the model to rely on multiple visual cues rather than a single discriminative region. MixUp creates convex combinations of training pairs, resulting in smoother decision boundaries. These techniques consistently improve accuracy by 1-3% on standard benchmarks.

Interview Questions on CNNs and Computer Vision

Technical interviews for computer vision roles test both theoretical understanding and practical implementation skills. Prepare for questions about architecture design, optimization challenges, and real-world deployment considerations.

Common Interview Topics

Interviewers frequently ask about receptive field calculations, the purpose of batch normalization, why convolutions work for images, and how to diagnose overfitting in vision models.

Q: What is the receptive field and why does it matter?

The receptive field is the region of the input image that influences a particular feature map location. Deeper layers have larger receptive fields, allowing them to capture broader context. For a network with n convolutional layers using 3x3 kernels, the receptive field grows as (2n + 1) × (2n + 1). Designing architectures requires balancing receptive field size against computational cost—dilated convolutions expand receptive fields without increasing parameters.

Q: Why use batch normalization in CNNs?

Batch normalization normalizes layer inputs to zero mean and unit variance, addressing internal covariate shift. Benefits include: faster convergence (higher learning rates become stable), regularization effect (reduces need for dropout), and gradient flow improvement (prevents vanishing gradients in deep networks). During inference, running statistics replace batch statistics for deterministic outputs.

Q: Explain the difference between same and valid padding.

Same padding adds zeros around the input to preserve spatial dimensions after convolution—a 32×32 input with a 3×3 kernel produces a 32×32 output. Valid padding applies no padding, reducing dimensions—the same operation produces a 30×30 output. Same padding is preferred in deep networks to avoid aggressive spatial reduction.

Q: How do you handle class imbalance in image classification?

Strategies include: weighted cross-entropy loss (higher weights for minority classes), oversampling minority classes during training, data augmentation focused on underrepresented classes, and focal loss which down-weights easy examples. For severe imbalance (>100:1), combine multiple techniques and consider metrics beyond accuracy—F1 score, precision-recall AUC, or confusion matrix analysis.

For more deep learning interview preparation, explore the CNN & Image Classification and Deep Learning Fundamentals question modules.

Deploying PyTorch Vision Models with TorchScript

Production deployment requires converting PyTorch models to optimized formats. TorchScript compiles models to a portable representation that runs without Python, enabling deployment in C++ applications, mobile devices, or serverless environments.

python

# deployment_export.py
import torch
from torchvision import models
from torchvision.models import ResNet18_Weights

def export_for_production(model: torch.nn.Module, save_path: str) -> None:
    """Export model to TorchScript for production deployment"""
    model.eval()  # Set to evaluation mode (disables dropout, uses running stats)
    
    # Create example input matching expected dimensions
    example_input = torch.randn(1, 3, 224, 224)
    
    # Method 1: Tracing (for models without control flow)
    traced_model = torch.jit.trace(model, example_input)
    traced_model.save(save_path.replace(".pt", "_traced.pt"))
    
    # Method 2: Scripting (handles if/else, loops)
    scripted_model = torch.jit.script(model)
    scripted_model.save(save_path.replace(".pt", "_scripted.pt"))
    
    # Verify outputs match
    with torch.no_grad():
        original_out = model(example_input)
        traced_out = traced_model(example_input)
        scripted_out = scripted_model(example_input)
    
    assert torch.allclose(original_out, traced_out, atol=1e-5)
    assert torch.allclose(original_out, scripted_out, atol=1e-5)
    print(f"Model exported successfully to {save_path}")

def optimize_for_inference(model_path: str) -> torch.jit.ScriptModule:
    """Load and optimize TorchScript model for inference"""
    model = torch.jit.load(model_path)
    
    # Optimize for inference (fuses operations, removes dropout)
    optimized = torch.jit.optimize_for_inference(model)
    
    return optimized

# Export ResNet18
model = models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
export_for_production(model, "resnet18_production.pt")

For maximum inference speed, consider ONNX export for cross-framework compatibility or TensorRT for NVIDIA GPU deployment. The PyTorch vs TensorFlow comparison covers deployment tradeoffs between frameworks.

Conclusion

Computer vision with PyTorch combines powerful abstractions with production-ready tools:

CNN architectures extract hierarchical features through convolution, pooling, and normalization layers
Transfer learning with pretrained models (ResNet, EfficientNet) achieves high accuracy on limited data
Data augmentation (MixUp, CutOut, geometric transforms) improves generalization and prevents overfitting
Technical interviews test receptive fields, batch normalization, padding strategies, and class imbalance handling
TorchScript enables production deployment without Python dependencies