手撸代码：从零开始的 AlexNet 图像分类（PyTorch框架）

摘要：

本文在 PyTorch 框架下搭建了 AlexNet ，并在 CIFAR10 上完成了图片分类。同时，更正了一些原论文中的小错误（如：输入图像尺寸）。由于 CIFAR10 没有验证集，本文将训练集的 10% 当作验证集。

完整代码已上传至 GitHub：https://github.com/TiezhuXing01/AlexNet_in_PyTorch

1. 引入库

1
2
3
4
5
6
import torch
import torch.nn as nn
import numpy as np
from torchvision import datasets
from torchvision import transforms
from torch.utils.data.sampler import SubsetRandomSampler

SubsetRandomSampler 是 PyTorch 中的一个采样器（sampler）。
具体可以看这篇文章：SubsetRandomSampler 是什么？

2. 选择设备

1
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

通常情况下，我们都会选择在GPU上训练网络模型，因为神经网络的训练需要大量的计算，而英伟达的GPU提供了CUDA（一个加速计算库）。但如果你的电脑显卡是AMD的，那么有很大概率不支持使用CUDA，此时只能用CPU训练。但在CPU上训练模型是十分缓慢的。如果你暂时没法换电脑，那我建议你去租一个服务器。或者使用阿里云、百度飞桨、谷歌Colab等平台。

3. 加载数据集

CIFAR-10 是一个经典的计算机视觉数据集，用于图像分类任务。它包含了来自 10 个不同类别的 60,000 张彩色图像，每个类别有 6,000 张图像。数据集被分为训练集和测试集，其中训练集包含 50,000 张图像，测试集包含 10,000 张图像。本文拿出训练集的 10% 作为验证集。

3.1 定义获取训练集和验证集的数据加载器

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def get_train_val_loader(data_dir, batch_size, augment,
                         random_seed, valid_size = 0.1, shuffle = True):
 
  # ------------- 设置图像变换 ------------- #
  # (1) 归一化
  normalize = transforms.Normalize(mean = [0.4914, 0.4822, 0.4465],
                                   std = [0.2023, 0.1994, 0.2010])
  # (2) 验证集图像变换
  val_transform = transforms.Compose([transforms.Resize(227),
                                   transforms.ToTensor(),
                                   normalize])
  # (3) 训练集是否数据增强
  if augment:
    train_transform = transforms.Compose([transforms.RandomHorizontalFlip(),
                                      transforms.Resize(227),
                                      transforms.ToTensor(),
                                      normalize])
  else:
    train_transform = transforms.Compose([transforms.Resize(227),
                                       transforms.ToTensor(),
                                       normalize])
  # ---------- 👆 数据变换设置完毕 -------------- #
 
  # 下载并加载训练集
  train_dataset = datasets.CIFAR10(root = data_dir,
                                  train = True,
                                  download = True,
                                  transform = train_transform)
  val_dataset = datasets.CIFAR10(root = data_dir,
                                  train = True,
                                  download = True,
                                  transform = val_transform)
 
  # ---------- 划分验证集和训练集 ----------  #
  # (1) 计算训练集图片数量
  num_train = len(train_dataset)
  # (2) 计算验证集数量，并向下取整
  num_val = np.floor(valid_size * num_train)
  # (3) 设置训练集和验证集的划分界限
  split = int(num_val)
  # (4) 生成一个列表索引，其内容为 0 ~ (num_train - 1) 的全部整数
  indices = list(range(num_train))    # 为数据"洗牌"做准备
  if shuffle:
    np.random.seed(random_seed) # 根据种子生成随机数
    np.random.shuffle(indices)  # 根据随机数打乱图片
  # (5) 划分验证集和训练集(根据索引列表 indices 和划分界限 split 划分)
  val_idx = indices[:split]     # 验证集索引列别
  train_idx = indices[split:]   # 训练集索引列表
  # (6) 根据验证集和训练集的索引列表采样数据
  train_sampler = SubsetRandomSampler(train_idx)
  val_sampler = SubsetRandomSampler(val_idx)
  # ---------- 👆 训练集和验证集划分完毕 ---------- #
 
  # 设置数据加载器
  train_loader = torch.utils.data.DataLoader(train_dataset,
                                             batch_size = batch_size,
                                             sampler = train_sampler)
  val_loader = torch.utils.data.DataLoader(val_dataset,
                                           batch_size = batch_size,
                                           sampler = val_sampler)
 
  return (train_loader, val_loader)

3.2 定义获取测试集的数据加载器

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def get_test_loader(data_dir, batch_size, shuffle = True):
 
  # ------------- 设置图像变换 ------------- #
  # 归一化
  normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                  std=[0.229, 0.224, 0.225],)
  # 图像变换
  test_transform = transforms.Compose([transforms.Resize(227),
                                transforms.ToTensor(),
                                normalize])
 
  # 下载并加载测试集
  test_dataset = datasets.CIFAR10(root = data_dir,
                                  train = False,
                                  download = True,
                                  transform = test_transform)
 
  # 加载测试数据
  test_loader = torch.utils.data.DataLoader(test_dataset,
                                            batch_size = batch_size,
                                            shuffle = shuffle)
 
  return test_loader

3.3 调用函数，加载数据

1
2
3
4
5
6
7
8
9
10
11
12
13
# 设置数据集下载路径
data_dir = "./data"
# 设置批尺寸
batch_size = 64
# 调用训练集和验证集的 DataLoader
train_loader, val_loader = get_train_val_loader(data_dir = data_dir,
                                                batch_size = batch_size,
                                                augment = True,
                                                random_seed = 1)
# 调用测试集的 DataLoader
test_loader = get_test_loader(data_dir = data_dir,
                              batch_size = batch_size,
                              shuffle = False)

4. 认识网络

AlexNet 结构图

AlexNet 是卷积神经网络的奠基之作，发表时间较早，所以今天看来有一些错误或者局限性：

(1) 在2012年 AlexNet 被提出时，GPU 的算力相对今天来说非常小。为了训练这个网络，他们把 AlexNet 从中间切开，使用多 GPU 训练，分别在 2 块 GTX 580 3G 上花了五、六天时间。

(2) 这种切割只是一种技术细节，以后几年的模型大多都不会这么切割，多 GPU 训练也不是通过切割模型来实现。所以本文的代码并没有切割网络，也没有用到多 GPU 训练。

(3) 论文中写到，输入图片的尺寸是 24∗24∗324*24*3 ，即 3 个通道的 24∗2424*24 图片。但事实上，这是作者的小错误。正确尺寸应为 27∗27∗327*27*3 .

如果按照论文说的输入尺寸做，会出现错误。可以来算一下， 卷积输出尺寸公式为：
outputsize=inputsize−kernelsize+2∗paddingstride+1output size=\frac{input size-kernel size+2*padding}{stride}+1
第一次卷积后：outputsize=224−11+2∗04+1=54.25output size=\frac{224-11+2*0}{4}+1=54.25 出现小数了。

在一块 GPU 上训练的正确 AlexNet 结构及其输入输出尺寸应为下表。
操作
名称		 输入
通道		 输出
通道		 核尺寸 步幅 填充 输出尺寸
输入: 227×227×3
conv1 3 96 11×11 4 0 55×55×96
pool 3 × 3 2 27×27×96
conv2 96 256 5 × 5 1 2 27×27×256
pool 3 × 3 2 13×13×256
conv3 256 384 3 × 3 1 1 13×13×384
conv4 384 384 3 × 3 1 1 13×13×384
conv5 384 256 3 × 3 1 1 13×13×256
pool 3 × 3 2 6×6×256
fc1 9216
fc2 4096
fc3 num_classes(=10)

5. 搭建网络

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class AlexNet(nn.Module):
    def __init__(self, num_classes=10):
        super(AlexNet, self).__init__()
        self.conv_block1 = nn.Sequential(
                        nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=0),
                        nn.BatchNorm2d(96),
                        nn.ReLU()
                        )
        self.conv_block2 = nn.Sequential(
                        nn.Conv2d(96, 256, kernel_size=5, stride=1, padding=2),
                        nn.BatchNorm2d(256),
                        nn.ReLU()
                        )
        self.conv_block3 = nn.Sequential(
                        nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1),
                        nn.BatchNorm2d(384),
                        nn.ReLU()
                        )
        self.conv_block4 = nn.Sequential(
                        nn.Conv2d(384, 384, kernel_size=3, stride=1, padding=1),
                        nn.BatchNorm2d(384),
                        nn.ReLU()
                        )
        self.conv_block5 = nn.Sequential(
                        nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1),
                        nn.BatchNorm2d(256),
                        nn.ReLU()
                        )
        self.pool = nn.MaxPool2d(kernel_size = 3, stride = 2)
        self.fc1 = nn.Sequential(nn.Dropout(0.5),nn.Linear(9216, 4096),nn.ReLU())
        self.fc2 = nn.Sequential(nn.Dropout(0.5),nn.Linear(4096, 4096),nn.ReLU())
        self.fc3 = nn.Sequential(nn.Linear(4096, num_classes))
 
    def forward(self, x):
        out = self.conv_block1(x)
        out = self.pool(out)
        out = self.conv_block2(out)
        out = self.pool(out)
        out = self.conv_block3(out)
        out = self.conv_block4(out)
        out = self.conv_block5(out)
        out = self.pool(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc1(out)
        out = self.fc2(out)
        out = self.fc3(out)
        return out
  • Dropout：让一部分神经元输出为0. 目前认为它是一个正则项，防止过拟合。由于目前很少使用全连接了，所以Dropout不是那么重要了。

6. 设置超参数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
num_classes = 10
epochs = 20
learning_rate = 0.005
 
model = AlexNet(num_classes).to(device)
 
# 设置损失函数
cost = nn.CrossEntropyLoss()
 
# 设置优化器
optimizer = torch.optim.SGD(model.parameters(),     # SGD 随机梯度下降
                            lr = learning_rate,
                            weight_decay = 0.005,   # 正则化项的权重是 0.005
                            momentum = 0.9)
 
# 一个 epoch 训练的总 step 数
train_step = len(train_loader)

  • len(train_loader)返回加载器中的批次数量。5000∗(1−0.1)64=703.125\frac{5000*(1-0.1)}{64}=703.125，结果出现小数，向上取整为704.

  • 随机梯度下降（SGD）现在是深度学习最主流的优化算法，因为其内部的噪声使模型有更好的泛化能力。

  • weight_decay：是 L2 的正则项。

  • 正则化：在机器学习中，正则化是一种用于防止过拟合的技术。限制模型的复杂性，防止模型对训练数据中的噪声过于敏感，从而提高其在未知数据上的泛化能力。在 AlexNet 提出那年（2012），人们普遍认为正则化对解决模型过拟合问题是很重要的。但在后期，这个观点被推翻了。取而代之的是，网络的设计对防止过拟合更重要的。

  • momentum：动量是一种优化算法中常用的技术，通常与随机梯度下降（SGD）结合使用，用于加速模型的训练过程。它的功能是 避免因下降曲线不平滑而落入局部最优解中。 动量的引入主要是为了解决随机梯度下降的一些问题，例如在梯度更新中存在的震荡和收敛速度慢的问题。动量算法引入了一个指数衰减的累积变量，用来持续跟踪梯度的历史信息。这个累积变量就是动量。动量在更新参数时不仅考虑当前梯度，还考虑了之前梯度的方向。这有助于平滑更新过程，减少参数更新的震荡，提高模型训练的稳定性和速度。

7. 训练和验证

训练需要 2 个循环的嵌套：外部循环用于循环 epoch ；内部循环用于循环每个 epoch 中的每个 batch 的图片，一个 batch 一步（step）。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
for epoch in range(epochs):
  for i, (images, labels) in enumerate(train_loader):
    images = images.to(device)
    labels = labels.to(device)
 
    outputs = model(images)
    loss = cost(outputs, labels)
 
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
 
  print("Epoch [{}/{}], Step [{}/{}], Loss:{:.4f}".format(
      epoch+1, epochs, i+1, train_step, loss.item()))   # loss是张量，需要.item()转为浮点型
 
  # 一个epoch完成之后，进入验证
  with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in val_loader:
      images = images.to(device)
      labels = labels.to(device)
      outputs = model(images)
      _, predicted = torch.max(outputs.data, 1)
      total += labels.size(0)
      correct += (predicted == labels).sum().item()
      # del images, labels, outputs   # 删除变量以释放内存
 
    # 输出验证结果
    print("Accuracy on validation: {} %".format(100 * (correct / total)))

  • 注意：enumerate(train_loader) 返回的是 (index, (images, labels)) 的元组，所以for i, images, labels in enumerate(train_loader):是错的，应该为for i, (images, labels) in enumerate(train_loader):. 相关文章：enumerate(train_loader) 返回什么？

  • torch.max(outputs.data, 1)返回每一行的最大值以及这些最大值所在的索引。第一个返回值（_）是最大值，第二个返回值（predicted）是最大值的索引。在这种情况下，我们只关心索引，因为它表示了模型的预测类别。

  • labels.size(0)返回的是当前批次中标签的数量。如何理解 labels.size(0)

  • del images, labels, outputs 是手动删除变量以释放内存。在Python中，这通常是不必要的，因为Python的垃圾回收器会自动处理不再使用的对象。

  • for images, labels in val_loader:不需要numerate(val_loader)是因为验证阶段不需要索引，看的是验证集整体的准确度。

训练结果如下图：

8.测试

1
2
3
4
5
6
7
8
9
10
11
12
13
14
with torch.no_grad():
  correct = 0
  total = 0
  for images, labels in test_loader:
    images = images.to(device)
    labels = labels.to(device)
    outputs = model(images)
    _, predicted = torch.max(outputs.data, 1)
    total += labels.size(0)
    correct += (predicted == labels).sum().item()
    # del images, labels, outputs   # 删除变量以释放内存
 
  # 输出测试结果
  print("Accuracy on test: {} %".format(100 * (correct / total)))

测试结果如下图：

以上，就是 AlexNet 实现图像分类的全部内容。