From Zero to PyTorch Hero: Building Image Classifiers#

Setting Up Your Environment#

We recommend using a virtual environment or Conda environment to manage dependencies:

1
# Create a new conda environment
2
conda create -n pytorch-env python=3.9
3

4
# Activate the environment
5
conda activate pytorch-env
6

7
# Install PyTorch (CPU-only, if you do not have a GPU)
8
conda install pytorch torchvision torchaudio cpuonly -c pytorch
9

10
# Or, if you have a CUDA-capable GPU, install the GPU version
11
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

Once installed, verify the PyTorch installation in a Python shell:

1
import torch
2
print(torch.__version__)

If you see a version number (e.g., 2.x.x), you’re all set!

Datasets and Transforms#

PyTorch’s torchvision module includes popular datasets like MNIST, CIFAR-10, and ImageNet, along with common transforms such as resizing, normalization, and random cropping. For instance:

1
import torchvision
2
import torchvision.transforms as transforms
3

4
transform = transforms.Compose([
5
    transforms.Resize((32, 32)),
6
    transforms.ToTensor(),
7
    transforms.Normalize((0.5,), (0.5,))  # Example for a single-channel image
8
])
9

10
train_dataset = torchvision.datasets.MNIST(
11
    root="data",
12
    train=True,
13
    transform=transform,
14
    download=True
15
)
16

17
test_dataset = torchvision.datasets.MNIST(
18
    root="data",
19
    train=False,
20
    transform=transform,
21
    download=True
22
)

In this example, we:

• Rezise every image to 32×32 pixels.
• Convert the PIL image to a PyTorch tensor with ToTensor().
• Normalize pixel values by subtracting the mean (0.5) and dividing by the standard deviation (0.5).

DataLoader and Batching#

The DataLoader class helps manage the dataset during training. It shuffles the data, loads it in batches, and can run multiple workers in parallel for faster loading:

1
from torch.utils.data import DataLoader
2

3
train_loader = DataLoader(
4
    dataset=train_dataset,
5
    batch_size=64,
6
    shuffle=True
7
)
8

9
test_loader = DataLoader(
10
    dataset=test_dataset,
11
    batch_size=64,
12
    shuffle=False
13
)

Now, train_loader and test_loader are ready to be iterated over in your training and evaluation loops.

Defining the Model#

We’ll start with a fully connected (feedforward) neural network for illustration:

1
import torch.nn as nn
2
import torch.nn.functional as F
3

4
class SimpleNN(nn.Module):
5
    def __init__(self, input_size=32*32, hidden_size=128, num_classes=10):
6
        super(SimpleNN, self).__init__()
7
        self.fc1 = nn.Linear(input_size, hidden_size)
8
        self.fc2 = nn.Linear(hidden_size, num_classes)
9

10
    def forward(self, x):
11
        # x shape is (batch_size, 1, 32, 32) for MNIST
12
        x = x.view(x.size(0), -1)  # Flatten
13
        x = F.relu(self.fc1(x))    # ReLU activation
14
        x = self.fc2(x)           # Output logits
15
        return x
16

17
model = SimpleNN()

Here:

• input_size is set to 32×32 = 1024 because our input images are reshaped to 32×32.
• We apply a hidden layer of 128 neurons, then an output layer with 10 neurons corresponding to the digits 0–9.
• We use the ReLU activation function on the hidden layer.

Loss Function and Optimizer#

We then specify the loss function (e.g., cross-entropy for classification) and an optimizer such as SGD:

1
import torch.optim as optim
2

3
criterion = nn.CrossEntropyLoss()
4
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

Training Loop#

A training loop typically involves:

1. Forward pass: Compute model output for the training batch.
2. Loss calculation: Compare predictions to labels with the loss function.
3. Backward pass: Compute gradients via loss.backward().
4. Update step: Use the optimizer to adjust weights.

1
num_epochs = 5
2

3
for epoch in range(num_epochs):
4
    model.train()  # Set model to train mode
5
    for images, labels in train_loader:
6
        # 1. Zero the gradient buffers
7
        optimizer.zero_grad()
8

9
        # 2. Forward pass
10
        outputs = model(images)
11

12
        # 3. Calculate loss
13
        loss = criterion(outputs, labels)
14

15
        # 4. Backward pass
16
        loss.backward()
17

18
        # 5. Update weights
19
        optimizer.step()
20

21
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")

As you train, you should see the loss decreasing over epochs. This indicates your model is learning to recognize MNIST digits.

Validation and Testing#

After each epoch (or after training completes), evaluate on the validation or test set to estimate model performance:

1
model.eval()  # Set model to eval mode
2
correct = 0
3
total = 0
4
with torch.no_grad():
5
    for images, labels in test_loader:
6
        outputs = model(images)
7
        _, predicted = torch.max(outputs.data, 1)
8
        total += labels.size(0)
9
        correct += (predicted == labels).sum().item()
10

11
accuracy = 100 * correct / total
12
print(f"Test Accuracy: {accuracy:.2f}%")

This snippet:

• Disables gradient computation with torch.no_grad().
• Computes the class with the highest score via torch.max.
• Calculates accuracy as (predicted == labels).sum().item() / total.

Why Convolutions?#

CNNs exploit local connectivity, enabling them to learn filters that activate when specific visual features (e.g., edges, corners, textures) appear in localized regions of the image. This yields better performance and fewer parameters compared to a fully connected architecture feeding on raw pixels.

Basic CNN Architecture#

A simple CNN structure for MNIST/CIFAR-10 might look like:

1
class SimpleCNN(nn.Module):
2
    def __init__(self, num_classes=10):
3
        super(SimpleCNN, self).__init__()
4
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1)
5
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
6
        self.conv2 = nn.Conv2d(16, 32, 3, 1, 1)
7
        self.fc1 = nn.Linear(32*8*8, num_classes)
8

9
    def forward(self, x):
10
        # x: (batch_size, 1, 32, 32)
11
        x = F.relu(self.conv1(x))  # (batch_size, 16, 32, 32)
12
        x = self.pool(x)           # (batch_size, 16, 16, 16)
13
        x = F.relu(self.conv2(x))  # (batch_size, 32, 16, 16)
14
        x = self.pool(x)           # (batch_size, 32, 8, 8)
15
        x = x.view(x.size(0), -1)  # Flatten: (batch_size, 32*8*8)
16
        x = self.fc1(x)            # (batch_size, num_classes)
17
        return x
18

19
model = SimpleCNN(num_classes=10)

Key points:

• Conv2d layers learn local features from 2D images.
• MaxPool2d reduces spatial dimensions.
• The final Linear layer maps features to class scores.

Try training this CNN with the same approach used for the fully connected network. You’ll likely see higher accuracy on image tasks.

Loading a Pretrained Model#

PyTorch’s torchvision.models contains many pretrained models. For instance, let’s take resnet18:

1
import torchvision.models as models
2

3
model = models.resnet18(pretrained=True)

This automatically loads weights from a ResNet-18 model trained on ImageNet.

Freezing and Unfreezing Layers#

Usually, you freeze the early layers (they learn universal features like edges) and only retrain the final layers to specialize in your dataset:

1
for param in model.parameters():
2
    param.requires_grad = False
3

4
# Replace the final classification layer
5
num_features = model.fc.in_features
6
model.fc = nn.Linear(num_features, 10)  # Suppose we have 10 classes

Now, only the new fc layer’s parameters will be learned. This method drastically reduces training time and data requirements.

Practical Tips#

• Small Dataset: Freeze more layers to avoid overfitting.
• Large Dataset: Unfreeze more layers for better capacity.
• Hyperparameters: Usually, a lower learning rate is used for fine-tuning pretrained networks.

Data Augmentation#

Image datasets can be small, so augmenting data can improve model generalization. Common techniques include random crops, flips, and color jitter:

1
transform_train = transforms.Compose([
2
    transforms.RandomResizedCrop(32),
3
    transforms.RandomHorizontalFlip(),
4
    transforms.ToTensor(),
5
    transforms.Normalize((0.5,), (0.5,))
6
])

This snippet randomly crops and horizontally flips images, then normalizes them. For test transforms, generally only resize and normalization are applied.

Learning Rate Schedulers#

Dynamically adjusting the learning rate can speed up convergence:

1
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
2

3
for epoch in range(num_epochs):
4
    # Training loop ...
5
    scheduler.step()

In this case, every 5 epochs, the learning rate is multiplied by 0.1. Other popular schedulers include ExponentialLR, ReduceLROnPlateau, and CosineAnnealingLR.

Regularization Methods#

Common ways to reduce overfitting include:

1. Dropout: Randomly zero some activations during training.
2. Weight Decay: A small penalty on large weights (L2 regularization).
3. Early Stopping: Stop training when validation loss stops improving.

1
class CNNWithDropout(nn.Module):
2
    def __init__(self, num_classes=10):
3
        super(CNNWithDropout, self).__init__()
4
        self.conv1 = nn.Conv2d(1, 16, 3, 1, 1)
5
        self.conv2 = nn.Conv2d(16, 32, 3, 1, 1)
6
        self.dropout = nn.Dropout(0.5)
7
        self.pool = nn.MaxPool2d(2, 2)
8
        self.fc1 = nn.Linear(32*8*8, num_classes)
9

10
    def forward(self, x):
11
        x = F.relu(self.conv1(x))
12
        x = self.pool(x)
13
        x = F.relu(self.conv2(x))
14
        x = self.pool(x)
15
        x = x.view(x.size(0), -1)
16
        x = self.dropout(x)
17
        x = self.fc1(x)
18
        return x

Batch Normalization and Layer Normalization#

Normalization layers stabilize training and boost performance. Batch Normalization normalizes each feature across the batch, while Layer Normalization normalizes across the features within a single sample:

1
self.bn1 = nn.BatchNorm2d(16)
2
...
3
x = F.relu(self.bn1(self.conv1(x)))

Such layers often allow for higher learning rates and faster convergence.

Quantization and Pruning#

For resource-constrained environments (e.g., mobile devices), you can consider:

• Quantization: Reducing float32 weights to int8 or int16.
• Pruning: Removing (zeroing out) weights with minimal impact on accuracy.

PyTorch offers built-in tooling for post-training and dynamic quantization, as well as for structured and unstructured pruning.

Deployment Strategies#

Common deployment paths include:

1. TorchScript: Convert your model to a serialized script or graph.
2. ONNX: Export to the Open Neural Network Exchange format for cross-platform deployment.
3. TensorRT (for NVIDIA GPUs): Highly optimized runtime for inference.

Example of exporting a PyTorch model to ONNX:

1
dummy_input = torch.randn(1, 1, 32, 32)
2
torch.onnx.export(model, dummy_input, "model.onnx", verbose=True)