使用 PyTorch 实现 AlexNet 进行 MNIST 图像分类

AlexNet 是一种经典的深度学习模型，它在 2012 年的 ImageNet 图像分类比赛中大放异彩，彻底改变了计算机视觉领域的格局。AlexNet 的核心创新包括使用深度卷积神经网络（CNN）来处理图像，并采用了多个先进的技术如 ReLU 激活函数、Dropout 正则化等。

本文将介绍如何使用 PyTorch 框架实现 AlexNet，并在 MNIST 数据集上进行训练。MNIST 是一个简单但经典的数据集，常用于初学者测试机器学习算法。

文末附完整项目。

一、AlexNet 网络结构

AlexNet 的结构大致可以分为两部分：特征提取部分（卷积层）和分类部分（全连接层）。下面是 AlexNet 的简要结构：

• 卷积层：五个卷积层用于提取特征。每个卷积层后面都有一个激活函数（ReLU）和一个池化层。
• 全连接层：三个全连接层，第一个和第二个全连接层后有 Dropout 层，防止过拟合。
• 输出层：使用 Softmax 激活函数输出 1000 个类别的概率。

二、使用 PyTorch 实现 AlexNet

训练部分

1. 导入必要的库

首先，我们需要导入一些必要的库，包括 PyTorch 和一些数据处理工具。

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision import datasets
from torch.cuda.amp import GradScaler, autocast  # 导入混合精度训练
from tqdm import tqdm
import multiprocessing
import matplotlib.pyplot as plt

2. 定义 AlexNet 模型

接下来，我们定义一个类 AlexNet，继承自 nn.Module，并在其中实现 AlexNet 的结构。

class AlexNet(nn.Module):
    def __init__(self):
        super(AlexNet, self).__init__()
        # 定义卷积层和池化层
        self.features = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=11, stride=4, padding=2),  # 第一个卷积层
            nn.ReLU(inplace=True),  # 激活函数
            nn.MaxPool2d(kernel_size=3, stride=2),  # 池化层
            nn.Conv2d(64, 192, kernel_size=5, padding=2),  # 第二个卷积层
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=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=3, stride=2)  # 池化层
        )
        # 定义全连接层
        self.classifier = nn.Sequential(
            nn.Dropout(),  # Dropout 层，防止过拟合
            nn.Linear(256 * 6 * 6, 4096),  # 第一个全连接层
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),  # 第二个全连接层
            nn.ReLU(inplace=True),
            nn.Linear(4096, 10)  # 输出层，10 个类别
        )

    def forward(self, x):
        # 前向传播
        x = self.features(x)  # 卷积层
        x = x.view(x.size(0), -1)  # 展平数据
        x = self.classifier(x)  # 全连接层
        return x

3. 数据预处理和加载

在进行训练之前，我们需要对 MNI