CNNs and Vision Transformers

Introduction to Deep Learning for Vision

Deep learning has revolutionized computer vision by automatically learning hierarchical feature representations from data. Unlike traditional methods that rely on hand-crafted features, deep neural networks learn features directly from raw pixels through multiple layers of abstraction.

Convolutional Neural Networks (CNNs)

Why Convolutions?

CNNs are specifically designed for processing grid-like data such as images. They leverage three key ideas:

• Local connectivity - Each neuron connects only to a small region of the input
• Parameter sharing - The same weights are used across different spatial locations
• Translation invariance - Features can be detected anywhere in the image

Convolutional Layer

import torch
import torch.nn as nn

# Define a convolutional layer
conv = nn.Conv2d(
    in_channels=3,      # RGB input
    out_channels=64,    # 64 feature maps
    kernel_size=3,      # 3x3 kernel
    stride=1,           # Slide by 1 pixel
    padding=1           # Keep spatial dimensions
)

# Forward pass
x = torch.randn(1, 3, 224, 224)  # Batch of 1, 3 channels, 224x224
output = conv(x)  # Shape: [1, 64, 224, 224]

# Mathematical operation:
# Output[b, c, h, w] = Σ Input[b, c', h+i, w+j] × Kernel[c, c', i, j]

Pooling Layers

Pooling reduces spatial dimensions while retaining important information:

# Max pooling - takes maximum value in each window
maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
x = torch.randn(1, 64, 224, 224)
pooled = maxpool(x)  # Shape: [1, 64, 112, 112]

# Average pooling
avgpool = nn.AvgPool2d(kernel_size=2, stride=2)

# Global average pooling - reduces to single value per channel
global_pool = nn.AdaptiveAvgPool2d((1, 1))
x = torch.randn(1, 512, 7, 7)
pooled = global_pool(x)  # Shape: [1, 512, 1, 1]

Activation Functions

# ReLU (Rectified Linear Unit) - most common
relu = nn.ReLU()
output = relu(x)  # max(0, x)

# LeakyReLU - allows small negative values
leaky_relu = nn.LeakyReLU(negative_slope=0.01)

# GELU (Gaussian Error Linear Unit) - used in transformers
gelu = nn.GELU()

# Sigmoid - outputs between 0 and 1
sigmoid = nn.Sigmoid()

# Tanh - outputs between -1 and 1
tanh = nn.Tanh()

Classic CNN Architectures

LeNet-5 (1998)

The pioneering CNN architecture for digit recognition.

class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

AlexNet (2012)

Sparked the deep learning revolution by winning ImageNet 2012 with 15.3% error (vs 26.2% for second place).

class AlexNet(nn.Module):
    def __init__(self, num_classes=1000):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )

VGGNet (2014)

Demonstrated that depth matters - achieved better results with 16-19 layers using only 3x3 convolutions.

# VGG-16 configuration
cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M',
       512, 512, 512, 'M', 512, 512, 512, 'M']

def make_layers(cfg):
    layers = []
    in_channels = 3
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)

ResNet (2015)

Introduced skip connections to solve vanishing gradient problem, enabling networks with 50-152+ layers.

class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super(ResidualBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels,
                               kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels,
                               kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels,
                         kernel_size=1, stride=stride),
                nn.BatchNorm2d(out_channels)
            )

    def forward(self, x):
        residual = x
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(residual)  # Skip connection
        out = F.relu(out)
        return out

# Key insight: F(x) + x is easier to optimize than F(x)

Inception (GoogLeNet) (2014)

Used parallel convolutions of different sizes to capture multi-scale features.

class InceptionModule(nn.Module):
    def __init__(self, in_channels, ch1x1, ch3x3_reduce, ch3x3,
                 ch5x5_reduce, ch5x5, pool_proj):
        super(InceptionModule, self).__init__()

        # 1x1 convolution branch
        self.branch1 = nn.Sequential(
            nn.Conv2d(in_channels, ch1x1, kernel_size=1),
            nn.ReLU(inplace=True)
        )

        # 1x1 -> 3x3 convolution branch
        self.branch2 = nn.Sequential(
            nn.Conv2d(in_channels, ch3x3_reduce, kernel_size=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(ch3x3_reduce, ch3x3, kernel_size=3, padding=1),
            nn.ReLU(inplace=True)
        )

        # 1x1 -> 5x5 convolution branch
        self.branch3 = nn.Sequential(
            nn.Conv2d(in_channels, ch5x5_reduce, kernel_size=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(ch5x5_reduce, ch5x5, kernel_size=5, padding=2),
            nn.ReLU(inplace=True)
        )

        # Max pooling -> 1x1 branch
        self.branch4 = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels, pool_proj, kernel_size=1),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        branch1 = self.branch1(x)
        branch2 = self.branch2(x)
        branch3 = self.branch3(x)
        branch4 = self.branch4(x)
        return torch.cat([branch1, branch2, branch3, branch4], 1)

EfficientNet (2019)

Systematically scales depth, width, and resolution for optimal efficiency.

# EfficientNet compound scaling
# depth = α^φ
# width = β^φ
# resolution = γ^φ
# subject to: α * β^2 * γ^2 ≈ 2

from torchvision.models import efficientnet_b0, efficientnet_b7

# Load pre-trained EfficientNet
model = efficientnet_b0(pretrained=True)
model.eval()

# EfficientNet-B0 to B7 variants
# B0: 5.3M params, 77.1% ImageNet accuracy
# B7: 66M params, 84.3% ImageNet accuracy

Vision Transformers (ViT)

The Attention Mechanism

Transformers use self-attention to model relationships between all parts of an image.

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        self.qkv = nn.Linear(embed_dim, embed_dim * 3)
        self.proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        B, N, C = x.shape

        # Generate Q, K, V
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        # Scaled dot-product attention
        attn = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5)
        attn = attn.softmax(dim=-1)

        # Apply attention to values
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        return x

Vision Transformer Architecture

class VisionTransformer(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_chans=3,
                 num_classes=1000, embed_dim=768, depth=12, num_heads=12):
        super().__init__()

        # Patch embedding
        self.patch_embed = nn.Conv2d(in_chans, embed_dim,
                                     kernel_size=patch_size,
                                     stride=patch_size)

        num_patches = (img_size // patch_size) ** 2

        # Learnable positional embeddings
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))

        # CLS token
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

        # Transformer encoder
        self.blocks = nn.ModuleList([
            TransformerBlock(embed_dim, num_heads)
            for _ in range(depth)
        ])

        # Classification head
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)

    def forward(self, x):
        B = x.shape[0]

        # Patch embedding
        x = self.patch_embed(x).flatten(2).transpose(1, 2)

        # Add CLS token
        cls_token = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_token, x), dim=1)

        # Add positional encoding
        x = x + self.pos_embed

        # Transformer blocks
        for block in self.blocks:
            x = block(x)

        x = self.norm(x)

        # Classify using CLS token
        return self.head(x[:, 0])

Hybrid Architectures

Combining CNNs and Transformers for best of both worlds:

• Swin Transformer - Hierarchical vision transformer with shifted windows
• ConvNeXt - Modernized CNN matching ViT performance
• CoAtNet - Combines convolution and attention
• MaxViT - Multi-axis attention for efficiency

Training Deep Networks

Loss Functions

# Cross-entropy loss for classification
criterion = nn.CrossEntropyLoss()
output = model(images)
loss = criterion(output, labels)

# Focal loss for imbalanced datasets
class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, inputs, targets):
        ce_loss = F.cross_entropy(inputs, targets, reduction='none')
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1-pt)**self.gamma * ce_loss
        return focal_loss.mean()

Optimization

# SGD with momentum
optimizer = torch.optim.SGD(model.parameters(), lr=0.01,
                           momentum=0.9, weight_decay=1e-4)

# Adam optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.001,
                            betas=(0.9, 0.999))

# AdamW (Adam with weight decay)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.001,
                             weight_decay=0.05)

# Learning rate scheduling
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=100
)

Data Augmentation

import torchvision.transforms as transforms

# Training augmentation
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2,
                          saturation=0.2, hue=0.1),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
])

# Validation (no augmentation)
val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
])

Regularization Techniques

# Dropout
dropout = nn.Dropout(p=0.5)

# Batch normalization
bn = nn.BatchNorm2d(num_features=64)

# Layer normalization (used in transformers)
ln = nn.LayerNorm(normalized_shape=768)

# Stochastic depth (drop path)
class DropPath(nn.Module):
    def __init__(self, drop_prob=0.0):
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        if not self.training or self.drop_prob == 0.:
            return x
        keep_prob = 1 - self.drop_prob
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        random_tensor = keep_prob + torch.rand(shape,
                                               dtype=x.dtype,
                                               device=x.device)
        random_tensor.floor_()
        return x.div(keep_prob) * random_tensor