精读 AlexNet 并使用 PyTorch 实现图像分类器

import torch
import torchvision
import torchvision.transforms as transforms

# 定义数据预处理：调整图片大小为 224x224（AlexNet 的输入尺寸），转换为 Tensor，并归一化
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 加载 CIFAR-10 数据集
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

# 定义 DataLoader，批量加载数据，设置 batch_size 为 32
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=32, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=32, shuffle=False)

import torch.nn as nn
import torch.nn.functional as F

class AlexNet(nn.Module):
    def __init__(self, num_classes=10):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            # 第一层卷积：输入 3 通道，输出 64 个特征图，卷积核 11x11，步长 4，填充 2
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            # 第二层卷积：输入 64，输出 192 个特征图，卷积核 5x5，填充 2
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            # 第三层卷积：输入 192，输出 384 个特征图，卷积核 3x3，填充 1
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            # 第四层卷积：输入 384，输出 256 个特征图，卷积核 3x3，填充 1
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            # 第五层卷积：输入 256，输出 256 个特征图，卷积核 3x3，填充 1
            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(),
            # 全连接层：输入由卷积层输出的特征图展平后，大小为 256*6*6
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes)
        )
    
    def forward(self, x):
        # 通过特征提取层
        x = self.features(x)
        # 将多维特征图展平为一维向量
        x = x.view(x.size(0), -1)
        # 通过分类器
        x = self.classifier(x)
        return x

# 创建模型实例并打印结构
model = AlexNet(num_classes=10)
print("模型结构：\n", model)

import torch.optim as optim

# 定义损失函数（交叉熵损失）和优化器（Adam）
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 设定训练轮数
num_epochs = 20

for epoch in range(num_epochs):
    running_loss = 0.0
    for i, (images, labels) in enumerate(train_loader):
        # 前向传播：将图片输入模型，获得预测结果
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        # 反向传播前先清空梯度
        optimizer.zero_grad()
        # 反向传播：自动计算梯度
        loss.backward()
        # 更新模型参数
        optimizer.step()
        
        running_loss += loss.item()
        if (i+1) % 100 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {running_loss/100:.4f}')
            running_loss = 0.0

print("训练完成！")

model.eval()  # 将模型设置为评估模式，关闭 Dropout 等训练特性
correct = 0
total = 0

with torch.no_grad():
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'测试集准确率：{100 * correct / total:.2f}%')

# 展示部分测试图片和预测结果
import matplotlib.pyplot as plt
import numpy as np

dataiter = iter(test_loader)
images, labels = next(dataiter)

def imshow(img):
    img = img / 2 + 0.5  # 反归一化
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

imshow(torchvision.utils.make_grid(images[:8]))

outputs = model(images[:8])
_, predicted = torch.max(outputs, 1)
print("预测结果：", predicted.numpy())

# 保存模型参数
torch.save(model.state_dict(), "alexnet_model.pth")
print("模型保存成功！")

# 加载模型参数
loaded_model = AlexNet(num_classes=10)
loaded_model.load_state_dict(torch.load("alexnet_model.pth"))
loaded_model.eval()
print("模型加载成功，开始推理...")

with torch.no_grad():
    outputs = loaded_model(images[:8])
    _, predicted = torch.max(outputs, 1)
    print("加载后预测结果：", predicted.numpy())