Writing VGG from Scratch in PyTorch

VGG (Visual Geometry Group) is a well-known convolutional neural network architecture that has achieved great success in image classification tasks. In this blog, we will learn how to implement the VGG architecture from scratch using PyTorch. We will cover the key components of the VGG network, including convolutional layers, pooling layers, and fully connected layers. By the end of this blog, you will have a solid understanding of how to build and train a VGG-like network in PyTorch.

Table of Contents#

  1. VGG Architecture Overview
  2. Implementing Convolutional Layers
  3. Implementing Pooling Layers
  4. Implementing Fully Connected Layers
  5. Putting It All Together
  6. Training the VGG Network
  7. Common Practices and Best Practices
  8. Example Usage
  9. References

1. VGG Architecture Overview#

The VGG architecture consists of multiple convolutional layers followed by pooling layers and finally fully connected layers. The key idea behind VGG is to use small-sized convolutional filters (e.g., 3x3) in multiple layers to capture hierarchical features in images. The network is typically deep, with a large number of convolutional layers.

The VGG-16 architecture, for example, has 16 weight layers (13 convolutional layers and 3 fully connected layers). The general structure is as follows:

  • Convolutional Layers: These layers perform feature extraction. They use convolutional operations with a specific kernel size (usually 3x3), stride, and padding.
  • Pooling Layers: Max pooling is commonly used to downsample the feature maps and reduce the spatial dimensions while retaining important features.
  • Fully Connected Layers: These layers are used for classification. They take the flattened output from the convolutional and pooling layers and map it to the number of classes.

2. Implementing Convolutional Layers#

In PyTorch, we can implement convolutional layers using the nn.Conv2d class. Here's an example of creating a simple convolutional layer:

import torch
import torch.nn as nn
 
# Create a convolutional layer
conv_layer = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1)

In the above code:

  • in_channels: The number of input channels (e.g., 3 for RGB images).
  • out_channels: The number of output channels (the number of filters in the convolutional layer).
  • kernel_size: The size of the convolutional kernel (3x3 in this case).
  • stride: The stride of the convolution operation.
  • padding: The amount of padding added to the input to keep the spatial dimensions the same (if padding = (kernel_size - 1) // 2 for a square kernel).

3. Implementing Pooling Layers#

Max pooling is implemented using the nn.MaxPool2d class in PyTorch. Here's an example:

# Create a max pooling layer
pool_layer = nn.MaxPool2d(kernel_size=2, stride=2)

The kernel_size and stride determine the size of the pooling window and how it moves across the feature map.

4. Implementing Fully Connected Layers#

Fully connected layers are implemented using the nn.Linear class. Here's an example:

# Create a fully connected layer
fc_layer = nn.Linear(in_features=25088, out_features=4096)

The in_features is the number of input features (usually the flattened size of the output from the convolutional and pooling layers), and out_features is the number of output features (the number of neurons in the fully connected layer).

5. Putting It All Together#

Now, let's build a simple VGG-like network. We'll use the VGG-11 architecture as an example (which has 11 weight layers).

class VGG11(nn.Module):
    def __init__(self, num_classes=1000):
        super(VGG11, self).__init__()
        self.features = nn.Sequential(
            # Block 1
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
 
            # Block 2
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
 
            # Block 3
            nn.Conv2d(128, 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=2, stride=2),
 
            # Block 4
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
 
            # Block 5
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
 
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, num_classes)
        )
 
    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

In the above code:

  • The features sequential container contains the convolutional and pooling layers.
  • The classifier sequential container contains the fully connected layers.
  • The forward method defines the forward pass of the network.

6. Training the VGG Network#

To train the VGG network, we need to define a loss function (e.g., nn.CrossEntropyLoss for classification) and an optimizer (e.g., torch.optim.SGD or torch.optim.Adam). Here's a simple training loop example:

# Assume we have a dataset and a data loader
# For simplicity, let's assume we have a train_loader
 
model = VGG11(num_classes=10)  # Adjust num_classes according to your dataset
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
 
num_epochs = 10
 
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for i, (images, labels) in enumerate(train_loader):
        optimizer.zero_grad()
 
        outputs = model(images)
        loss = criterion(outputs, labels)
 
        loss.backward()
        optimizer.step()
 
        running_loss += loss.item()
 
    print(f'Epoch {epoch + 1}, Loss: {running_loss / len(train_loader)}')

7. Common Practices and Best Practices#

  • Data Preprocessing: Normalize the input images (e.g., subtract mean and divide by standard deviation) to improve training.
  • Weight Initialization: Use proper weight initialization techniques (e.g., nn.init.xavier_uniform_ for convolutional and fully connected layers).
  • Regularization: Use techniques like dropout (as in the classifier part of our VGG implementation) to prevent overfitting.
  • Learning Rate Scheduling: Use learning rate schedulers (e.g., torch.optim.lr_scheduler.StepLR) to adjust the learning rate during training.

8. Example Usage#

Here's how you can use the trained VGG network for inference:

model.eval()
with torch.no_grad():
    test_image = torch.randn(1, 3, 224, 224)  # Assume input size is 224x224 for a typical VGG input
    output = model(test_image)
    _, predicted = torch.max(output.data, 1)
    print(f'Predicted class: {predicted.item()}')

9. References#

This blog has provided a detailed guide on writing VGG from scratch in PyTorch. By following the steps and best practices, you can build and train your own VGG-like networks for various image classification tasks.

2026-05