Python手写CNN:50行代码,从网络定义到训练评估

本文用 PyTorch 搭建一个三层卷积网络，加入 BatchNorm 加速收敛。代码完整可运行，重点看三件事：网络结构、训练循环、评估方法。

网络定义

import torchimport torch.nn as nnclass CNN(nn.Module):    def __init__(self):        super().__init__()        self.conv = nn.Sequential(            nn.Conv2d(1, 32, 3, padding=1),   # 1→32通道，3×3卷积，保持尺寸            nn.BatchNorm2d(32),                # 批归一化：加速训练、稳定梯度            nn.ReLU(),            nn.MaxPool2d(2),                   # 尺寸减半            nn.Conv2d(32, 64, 3, padding=1),  # 32→64通道            nn.BatchNorm2d(64),            nn.ReLU(),            nn.MaxPool2d(2),            nn.Conv2d(64, 128, 3, padding=1), # 64→128通道            nn.BatchNorm2d(128),            nn.ReLU(),            nn.AdaptiveAvgPool2d(1),           # 全局平均池化：输出1×1        )        self.fc = nn.Sequential(            nn.Flatten(),            nn.Linear(128, 64), nn.ReLU(),            nn.Dropout(0.3),                   # 随机丢弃30%神经元，防过拟合            nn.Linear(64, 10),                 # 10分类        )    def forward(self, x):        return self.fc(self.conv(x))

逐行看懂：BatchNorm2d 对每个通道做归一化，让数据分布稳定，训练更快。AdaptiveAvgPool2d(1) 把任意尺寸特征图压缩成 1×1×128，替代传统的 Flatten+全连接，参数更少、更鲁棒。

训练循环

model = CNN()loss_fn = nn.CrossEntropyLoss()opt = torch.optim.Adam(model.parameters(), lr=0.001)for epoch in range(5):    for x, y in train_loader:        opt.zero_grad()        loss = loss_fn(model(x), y)        loss.backward()        opt.step()    print(f"epoch {epoch+1} done")

Adam 优化器自动调整学习率，比 SGD 更省心。zero_grad→forward→backward→step 四行循环是深度学习训练的全部精髓。

评估

correct = total = 0model.eval()                                    # 切换到评估模式（关闭Dropout）with torch.no_grad():    for x, y in test_loader:        correct += (model(x).argmax(1) == y).sum().item()        total += y.size(0)print(f"准确率: {100*correct/total:.2f}%")

model.eval() 是关键：评估时必须关闭 Dropout 和 BatchNorm 的统计更新，否则结果会抖动。

三个新手常见错误

1. 1. 忘记 model.eval()：测试准确率忽高忽低，就是因为 Dropout 还在随机丢弃。
2. 2. 图片没归一化：像素值 0~255 会让梯度爆炸，训练前务必除以 255。
3. 3. 混淆 nn.CrossEntropyLoss 和 nn.NLLLoss：CrossEntropyLoss 内部已经做了 Softmax，网络最后一层不需要再加 Softmax。

下一篇，用这个网络跑真实数据集 CIFAR-10，看实战效果