AlexNet训练Fashion-MNIST数据集

import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torchvision import transforms

# AlexNet模型
net = nn.Sequential(
    nn.Conv2d(in_channels=1, out_channels=96, kernel_size=11, stride=4, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2),

    nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, padding=2),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2),

    nn.Conv2d(256, 384, kernel_size=3, padding=1),
    nn.ReLU(),

    nn.Conv2d(384, 384, kernel_size=3, padding=1),
    nn.ReLU(),

    nn.Conv2d(384, 256, kernel_size=3, padding=1),
    nn.ReLU(),

    nn.MaxPool2d(kernel_size=3, stride=2),

    nn.Flatten(),

    # 这里使用dropout层来减轻过拟合
    nn.Linear(in_features=6400, out_features=4096),
    nn.ReLU(),

    nn.Dropout(p=0.5),

    nn.Linear(4096, 4096),
    nn.ReLU(),

    nn.Dropout(p=0.5),

    nn.Linear(4096, 10)
)

# 将一个大小为 224*224 的单通道（黑白）图像通过LeNet
X = torch.rand(size=(1, 1, 224, 224), dtype=torch.float32)
# 在每一层打印输出的形状
print("AlexNet神经网络结构：")
for layer in net:
    X = layer(X)
    print(layer.__class__.__name__, 'output shape: \t', X.shape)

# GPU训练
device = torch.device("cuda")
net = net.to(device)

# 使用 FashionMNIST 服装数据集（10分类）
# 训练模型

# 原文中AlexNet是在ImageNet上进行训练的
# Fashion-MNIST图像的分辨率28*28（像素）低于ImageNet图像，转换为224*224
# （通常来讲这不是一个明智的做法，但在这里这样做是为了有效使用AlexNet架构）

transform = torchvision.transforms.Compose([
    transforms.Resize(size=(224, 224)),
    transforms.ToTensor()
])

train_data = torchvision.datasets.FashionMNIST(root="../../data", train=True,
                                               transform=transform,
                                               download=True)
test_data = torchvision.datasets.FashionMNIST(root="../../data", train=False,
                                              transform=transform,
                                              download=True)

print(f"len(train_data)={len(train_data)}")
print(f"len(test_data)={len(test_data)}")

train_loader = DataLoader(dataset=train_data, batch_size=128, shuffle=True)
test_loader = DataLoader(dataset=test_data, batch_size=128, shuffle=True)
print(f"len(train_loader)={len(train_loader)}")
print(f"len(test_loader)={len(test_loader)}")

learning_rate = 0.01
optimizer = torch.optim.Adam(params=net.parameters(), lr=learning_rate)
loss_func = nn.CrossEntropyLoss()
loss_func = loss_func.to(device)
epochs = 5
total_train_step = 0

for epoch in range(epochs):
    net.train()
    print(f"第{epoch + 1}次训练开始:")

    for data in train_loader:
        imgs, labels = data
        imgs = imgs.to(device)
        labels = labels.to(device)
        outputs = net(imgs)
        loss = loss_func(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step += 1
        if total_train_step % 100 == 0:
            print(f"训练次数：{total_train_step},loss:{loss.item()}")

net.eval()
correct = 0
total = 0
with torch.no_grad():
    for data in test_loader:
        imgs, labels = data
        imgs = imgs.to(device)
        labels = labels.to(device)
        outputs = net(imgs)
        _, predicted = torch.max(outputs.data, 1)  # 找到模型输出中概率最高的类别的索引，即预测的类别
        total += labels.size(0)  # 将当前批次的样本数量加到total变量中
        correct += (predicted == labels).sum().item()  # 计算当前批次中预测正确的样本数量,并加到correct变量中

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