这篇把 AlexNet 用到五分类花卉数据集上。流程和前面的 LeNet Demo 相同，只是多了数据集划分，模型也更深：准备 ImageFolder 目录、定义网络、训练并保存最优权重，最后预测单张图片。

我在全连接层使用 2048 个节点，而不是原论文的 4096 个节点。这样能减少本地显存占用，所以它是适合练习的 AlexNet 变体，不是逐参数复刻论文。

准备数据集

数据集下载地址：http://download.tensorflow.org/example_images/flower_photos.tgz

原始目录中每个子目录代表一个类别：

1
2
3
4
5
6
7
flower_data/
  flower_photos/
    daisy/
    dandelion/
    roses/
    sunflowers/
    tulips/

下面的脚本把 10% 图片随机划到验证集。它会重新创建 train 和 val 目录，运行前不要在这两个目录里放其他文件。

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
import random
import shutil
from pathlib import Path


def recreate_dir(path: Path):
    if path.exists():
        shutil.rmtree(path)
    path.mkdir(parents=True)


def main():
    random.seed(0)
    split_rate = 0.1
    data_root = Path("flower_data")
    source_root = data_root / "flower_photos"
    train_root = data_root / "train"
    val_root = data_root / "val"

    if not source_root.exists():
        raise FileNotFoundError(source_root)

    classes = [path for path in source_root.iterdir() if path.is_dir()]
    recreate_dir(train_root)
    recreate_dir(val_root)

    for class_dir in classes:
        images = [path for path in class_dir.iterdir() if path.is_file()]
        val_images = set(random.sample(images, k=int(len(images) * split_rate)))
        (train_root / class_dir.name).mkdir()
        (val_root / class_dir.name).mkdir()

        for image in images:
            target_root = val_root if image in val_images else train_root
            shutil.copy2(image, target_root / class_dir.name / image.name)

        print(f"{class_dir.name}: train={len(images) - len(val_images)}, val={len(val_images)}")


if __name__ == "__main__":
    main()

定义 AlexNet：model.py

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
import torch
import torch.nn as nn


class AlexNet(nn.Module):
    def __init__(self, num_classes=1000):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(96, 256, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(256, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(256 * 6 * 6, 2048),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(2048, 2048),
            nn.ReLU(inplace=True),
            nn.Linear(2048, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, start_dim=1)
        return self.classifier(x)

AdaptiveAvgPool2d((6, 6)) 明确了全连接层需要的输入尺寸。这样即使前面的输入尺寸略有变化，只要特征图不小于目标尺寸，分类器仍能收到固定长度的向量。

训练：train.py

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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import json
from pathlib import Path

import torch
import torch.nn as nn
from torch.optim import Adam
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

from model import AlexNet


def evaluate(model, loader, device):
    model.eval()
    correct = total = 0
    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            predictions = model(images).argmax(dim=1)
            total += labels.size(0)
            correct += (predictions == labels).sum().item()
    return correct / total


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    data_root = Path("flower_data")

    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])
    val_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])

    train_set = datasets.ImageFolder(data_root / "train", transform=train_transform)
    val_set = datasets.ImageFolder(data_root / "val", transform=val_transform)
    train_loader = DataLoader(train_set, batch_size=32, shuffle=True, num_workers=0)
    val_loader = DataLoader(val_set, batch_size=32, shuffle=False, num_workers=0)

    class_indices = {index: name for name, index in train_set.class_to_idx.items()}
    Path("class_indices.json").write_text(
        json.dumps(class_indices, ensure_ascii=False, indent=2), encoding="utf-8"
    )

    model = AlexNet(num_classes=len(class_indices)).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = Adam(model.parameters(), lr=0.0002)
    best_accuracy = 0.0

    for epoch in range(10):
        model.train()
        running_loss = 0.0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            loss = criterion(model(images), labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        accuracy = evaluate(model, val_loader, device)
        print(
            f"epoch {epoch + 1}: "
            f"loss={running_loss / len(train_loader):.4f}, "
            f"val_acc={accuracy:.4f}"
        )
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            torch.save(model.state_dict(), "alexnet-flowers.pth")


if __name__ == "__main__":
    main()

ImageFolder 会按目录名排序并生成类别编号，因此预测脚本不能自己手写另一套编号。训练时把映射保存成 JSON，后面直接读取最稳妥。

预测：predict.py

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
import json

import torch
from PIL import Image
from torchvision import transforms

from model import AlexNet


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])

with open("class_indices.json", encoding="utf-8") as file:
    class_indices = json.load(file)

image = transform(Image.open("test.jpg").convert("RGB")).unsqueeze(0).to(device)
model = AlexNet(num_classes=len(class_indices)).to(device)
model.load_state_dict(torch.load("alexnet-flowers.pth", map_location=device))
model.eval()

with torch.no_grad():
    probabilities = torch.softmax(model(image), dim=1)[0]

index = probabilities.argmax().item()
print(class_indices[str(index)], probabilities[index].item())

我遇到过的问题

DataLoader 迭代器没有 next 方法

旧教程里经常出现：

1
images, labels = data_iter.next()

Python 3 中直接使用内置函数：

1
images, labels = next(data_iter)

CUDA 显存不足

先减小 batch size，例如从 32 改成 16 或 8。单个 batch 仍然放不下时，再缩小全连接层、输入尺寸或模型通道数。减少 epoch 只会缩短总训练时间，不能解决一次前向传播造成的显存不足。

这次实践让我真正理解了 AlexNet 结构图之外的部分：数据目录、类别映射、训练和预测预处理是否一致，往往比某一层参数更容易让结果出错。