A Gentle Introduction to Deep Learning - AlexNet Code Implementation

published: and updated:

Deep Learning Deep Learning , CNN , AlexNet

Language / 语言

English (current) | 简体中文

Preface#

In this section, I will implement the AlexNet network construction and training process from a code perspective, guiding readers through the entire process from data collection to final model performance evaluation in deep learning.

I will build an AlexNet from scratch and train it on the CIFAR-10 dataset, then perform inference using images from the web.

The network will be built on PyTorch. Let’s start with the core network structure.

TIP

This assumes readers have some foundation in Python syntax and object-oriented programming. If you encounter any confusing methods or syntax, please refer to relevant resources promptly.

Step 1. Data Preprocessing#

Before building the network, we first need to process the input data. For the CIFAR-10 dataset, image dimensions are 32×32, which differs from the original AlexNet’s 224×224 input, so we need to adjust the network structure.

python
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image

NUM_CLASSES = 10

# Data preprocessing pipeline
tran = transforms.Compose([
    transforms.Resize((32,32), interpolation=Image.BICUBIC),
    transforms.ToTensor(), # Convert to tensor, high-dimensional array
    transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.201])
])

The data preprocessing includes three steps:

  1. Size adjustment: Resize images to 32×32 pixels
  2. Tensor conversion: Convert PIL images to PyTorch tensors
  3. Normalization: Normalize using CIFAR-10 dataset’s mean and standard deviation
TIP

Normalization is an important preprocessing step in deep learning that can accelerate model convergence and improve training stability. The values used here are statistical values for the CIFAR-10 dataset across RGB channels.

Step 2. AlexNet Network Construction#

To build our own network, we first need to import the torch.nn package, which is a core module containing neural network methods and classes.

python
class AlexNet(nn.Module):
    def __init__(self, num_class=NUM_CLASSES):
        super(AlexNet, self).__init__()

For our network, we only need the __init__ method to initialize the network structure and the forward method to define forward propagation. You may have seen other versions of AlexNet elsewhere containing hundreds of lines of code, most of which are for data processing or adding data statistics and storage functions. After understanding the core code, you can freely complete the remaining parts.

Feature Extraction Layers#

python
self.features = nn.Sequential(
    nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1),
    nn.ReLU(inplace=True), 
    nn.MaxPool2d(kernel_size=2),
    nn.Conv2d(64, 192, kernel_size=3, padding=1),
    nn.ReLU(inplace=True),
    nn.MaxPool2d(kernel_size=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=2),
)

Readers might find these various parameters complex. Here are the detailed parameters for each layer targeting 32×32 input:

LayerInput SizeOutput SizeKernel SizeStridePadding
Input32×32×3----
Conv-132×32×316×16×643×321
Pool-116×16×648×8×642×220
Conv-28×8×648×8×1923×311
Pool-28×8×1924×4×1922×220
Conv-34×4×1924×4×3843×311
Conv-44×4×3844×4×2563×311
Conv-54×4×2564×4×2563×311
Pool-34×4×2562×2×2562×220

ReLU activation is applied after each convolution.

TIP

Here I use ReLU instead of Sigmoid to enhance training accuracy and prevent the vanishing gradient problem mentioned earlier.

Classifier Layers#

As mentioned before, fully connected layers receive one-dimensional vectors, so we use the nn.Sequential method to build the classifier:

python
self.classifier = nn.Sequential(
    nn.Dropout(),
    nn.Linear(256*2*2, 4096),
    nn.ReLU(inplace=True),
    nn.Dropout(),
    nn.Linear(4096, 4096),
    nn.ReLU(inplace=True),
    nn.Linear(4096, num_class),
)

The calculation 256*2*2 comes from the output dimensions of the last pooling layer: 2×2×256.

Forward Propagation#

python
def forward(self, x):
    x = self.features(x)
    x = x.view(x.size(0), 256*2*2)
    x = self.classifier(x)
    return x

The view method here serves as a bridge. Input images, after passing through multiple convolutional layers, are output as tensors, but fully connected layers must receive one-dimensional vectors, so flattening is necessary. The view method parameters are:

python
view(batch_size, flattened_features)

Therefore, flatten the tensor while keeping batch_size unchanged.

Step 3. Training Process#

The training process is a cycle of forward and backward propagation repeatedly, but in practice, it’s much more complex.

Here we need to introduce several important concepts:

Learning Rate Adjustment Strategies#

  1. Learning Rate: Learning rate controls the speed of loss function changes. Too small a learning rate will cause the model to learn too slowly and get stuck in local optima. Too aggressive a learning rate will cause parameters to oscillate around the optimal solution, making convergence difficult.

  2. Warmup: The warmup process is necessary in early training, gradually increasing the learning rate from a very small value to the target value. Generally, linear learning rate warmup is used:

    When training steps are less than warmup steps, learning rate is updated as:

    lr=start_lr+(end_lrstart_lr)×epochwarmup_stepslr = start\_lr + (end\_lr - start\_lr) \times \frac{epoch}{warmup\_steps}

    When training steps are greater than or equal to warmup steps:

    lr=learning_ratelr = learning\_rate
  3. Learning Rate Scheduler: In actual training, using a fixed learning rate is unreliable and will prevent accuracy from improving further. Therefore, we can adjust the learning rate downward after specific epochs.

Trainer Implementation#

I implemented a complete trainer class with the following core functions:

python
class Trainer(object):
    def __init__(self, model_name, model, lr, train_on_gpu=False):
        self.model = model
        self.lr = lr
        self.model_name = model_name
        self.train_on_gpu = train_on_gpu
        self.best_acc = 0
        self.best_epoch = 0
        
        if self.train_on_gpu:
            self.model = self.model.cuda()
        
        # Optimizer setup
        self.optimizer = optim.SGD(
            self.model.parameters(),
            self.lr,
            momentum=0.9,
            weight_decay=5e-4
        )
        
        # Learning rate scheduler
        self.scheduler = MultiStepLR(
            self.optimizer, 
            milestones=[10, 20, 50, 100, 180], 
            gamma=0.1
        )

Warmup Implementation#

python
def warmup_learning_rate(self, init_lr, no_of_steps, epoch, len_epoch):
    """Warmup learning rate for first 5 epochs"""
    factor = no_of_steps // 30
    lr = init_lr * (0.1**factor)
    # Warmup calculation
    lr = lr * float(1 + epoch + no_of_steps * len_epoch) / (5. * len_epoch)
    return lr

Training Loop#

python
def train(self, epoch, no_of_steps, trainloader):
    self.model.train()
    train_loss, correct, total = 0, 0, 0
    
    # Use warmup for first 5 epochs, then use scheduler
    if epoch < 5:
        lr = self.warmup_learning_rate(self.lr, no_of_steps, epoch, len(trainloader))
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr
    elif epoch == 5:
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = self.lr
    
    criterion = nn.CrossEntropyLoss()
    
    for idx, (inputs, targets) in enumerate(trainloader):
        if self.train_on_gpu:
            inputs, targets = inputs.cuda(), targets.cuda()
            
        self.model.zero_grad()
        outputs = self.model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        self.optimizer.step()
        
        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += (targets == predicted).sum().item()
    
    if epoch >= 5:
        self.scheduler.step()
    
    return 100. * correct / total

Step 4. Model Testing and Inference#

After training, we need to test model performance and perform actual inference.

Test Function#

python
def test():
    net = AlexNet().cuda()
    model_path = os.path.join("weights", "alexnet.pt")
    
    # Load trained model
    checkpoint = torch.load(model_path)
    net.load_state_dict(checkpoint['net'])
    
    # Load test image
    test_image = os.path.join('test.jpg')
    img = Image.open(test_image)
    img_tensor = tran(img)  # Apply preprocessing
    
    # Add batch dimension and move to GPU
    input_tensor = img_tensor.unsqueeze_(0).cuda()
    
    # Model inference
    y = net(input_tensor)
    
    # Calculate probability distribution
    percentage = torch.softmax(y[0], dim=0) * 100
    cl_fp32, index_fp32 = torch.max(percentage, 0)
    
    # CIFAR-10 classes
    classes = ['plane', 'car', 'bird', 'cat', 'deer', 
               'dog', 'frog', 'horse', 'ship', 'truck']
    
    # Add prediction result to image
    font = ImageFont.truetype('LiberationSans-Regular.ttf', 30)
    draw = ImageDraw.Draw(img)
    text = str(classes[index_fp32]) + ' (' + '{:.2f}'.format(cl_fp32.item()) + '%)'
    draw.text((0, 0), text, font=font, fill="#ff00ff")
    
    img.save(test_image, 'jpeg')
    print(f'Prediction: {classes[index_fp32]} ({cl_fp32.item():.2f}%)')

Step 5. Training Monitoring and Visualization#

To better understand the training process, the trainer also includes complete monitoring and visualization functions:

Performance Curve Plotting#

python
def plot_accuracy_curves(self):
    import matplotlib.pyplot as plt
    
    epochs = range(1, len(self.train_acc_history) + 1)
    
    plt.figure(figsize=(12, 8))
    plt.plot(epochs, self.train_acc_history, 'b-', label='Training Accuracy', linewidth=2)
    plt.plot(epochs, self.test_acc_history, 'r-', label='Test Accuracy', linewidth=2)
    
    # Mark best test accuracy
    best_epoch = self.best_epoch + 1
    best_acc = self.best_acc
    plt.plot(best_epoch, best_acc, 'ro', markersize=10,
             label=f'Best Test Acc: {best_acc:.2f}% (Epoch {best_epoch})')
    
    # Mark learning rate adjustment points
    milestones = [10, 20, 50, 100, 180]
    for ms in milestones:
        if ms <= len(epochs):
            plt.axvline(x=ms, color='gray', linestyle='--', alpha=0.7)
    
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.title(f'AlexNet Training on CIFAR-10\\nBest Test Acc: {best_acc:.2f}%')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.savefig('accuracy_curves.png', dpi=300)

Experimental Results and Analysis#

Through training on the CIFAR-10 dataset, we can observe:

  1. Importance of Warmup: The 5-epoch warmup process ensures training stability
  2. Learning Rate Scheduling: Reducing learning rate at specific epochs helps the model converge to better local optima
  3. Overfitting Control: Dropout layers effectively prevent overfitting

Typical training curves will show:

  • Training accuracy steadily increasing
  • Test accuracy rising rapidly initially, then stabilizing
  • Noticeable performance improvements at learning rate adjustment points

Summary#

Through this section’s learning, readers should have mastered:

  1. Network Implementation: How to build AlexNet using PyTorch
  2. Training Pipeline: Including data preprocessing, model training, performance monitoring
  3. Optimization Techniques: Learning rate scheduling, warmup, dropout, and other techniques
  4. Model Evaluation: How to evaluate and visualize model performance

This complete implementation demonstrates the typical workflow of deep learning projects, covering every aspect from data preparation to model deployment. In the next section, we will explore more advanced network architectures and optimization techniques.

TIP

Practice Recommendation: I recommend readers run this code in their own environment to deepen understanding through hands-on experience. Adjust hyperparameters and observe how different settings affect model performance - this is the best way to learn deep learning.