A Gentle Introduction to Deep Learning - AlexNet Code Implementation
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.
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.
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:
- Size adjustment: Resize images to 32×32 pixels
- Tensor conversion: Convert PIL images to PyTorch tensors
- Normalization: Normalize using CIFAR-10 dataset’s mean and standard deviation
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.
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
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:
| Layer | Input Size | Output Size | Kernel Size | Stride | Padding |
|---|---|---|---|---|---|
| Input | 32×32×3 | - | - | - | - |
| Conv-1 | 32×32×3 | 16×16×64 | 3×3 | 2 | 1 |
| Pool-1 | 16×16×64 | 8×8×64 | 2×2 | 2 | 0 |
| Conv-2 | 8×8×64 | 8×8×192 | 3×3 | 1 | 1 |
| Pool-2 | 8×8×192 | 4×4×192 | 2×2 | 2 | 0 |
| Conv-3 | 4×4×192 | 4×4×384 | 3×3 | 1 | 1 |
| Conv-4 | 4×4×384 | 4×4×256 | 3×3 | 1 | 1 |
| Conv-5 | 4×4×256 | 4×4×256 | 3×3 | 1 | 1 |
| Pool-3 | 4×4×256 | 2×2×256 | 2×2 | 2 | 0 |
ReLU activation is applied after each convolution.
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:
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
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), 256*2*2)
x = self.classifier(x)
return xThe 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:
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
-
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.
-
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:
When training steps are greater than or equal to warmup steps:
-
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:
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
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 lrTraining Loop
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 / totalStep 4. Model Testing and Inference
After training, we need to test model performance and perform actual inference.
Test Function
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
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:
- Importance of Warmup: The 5-epoch warmup process ensures training stability
- Learning Rate Scheduling: Reducing learning rate at specific epochs helps the model converge to better local optima
- 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:
- Network Implementation: How to build AlexNet using PyTorch
- Training Pipeline: Including data preprocessing, model training, performance monitoring
- Optimization Techniques: Learning rate scheduling, warmup, dropout, and other techniques
- 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.
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.