PyTorch实战：从零构建CIFAR-10图像分类器（含训练、测试与验证集全流程解析）

1. 环境准备与数据加载

第一次用PyTorch做图像分类时，我被CIFAR-10数据集的可爱小图片吸引了——32x32像素的飞机、青蛙、汽车，像极了小时候玩的贴纸。这个经典数据集包含6万张图片，分为10个类别，正好适合练手。下面分享我摸索出来的完整操作流程。

安装PyTorch时有个小技巧：直接去官网复制对应CUDA版本的命令。比如我的RTX 3060笔记本跑这段代码：

bash复制pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

加载数据集时要注意三个关键参数：

python复制train_data = torchvision.datasets.CIFAR10(
    root='./data',  # 数据存储路径
    train=True,     # 训练集模式
    transform=torchvision.transforms.ToTensor(),  # 自动转为张量
    download=True   # 自动下载
)

这里有个实际项目中的经验：transform参数可以玩出花样。比如我常加个随机水平翻转增强数据：

python复制transform = torchvision.transforms.Compose([
    torchvision.transforms.RandomHorizontalFlip(p=0.5),
    torchvision.transforms.ToTensor()
])

DataLoader的batch_size设置也有讲究。我的1660Ti显卡跑64很流畅，但之前用MX450时就得降到32。打印数据集长度能快速检查是否加载成功：

python复制print(f"训练集样本数: {len(train_data)}")  # 输出50000
print(f"测试集样本数: {len(test_data)}")   # 输出10000

2. 构建CNN网络

刚开始学CNN时，我总纠结每层的参数该怎么设。后来发现CIFAR-10这种小图片，用三层卷积+池化组合就够用。这个结构是我参考多个开源项目调试出来的：

python复制class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            # 卷积层1: 3通道→32通道，5x5卷积核
            nn.Conv2d(3, 32, kernel_size=5, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(2),
            # 卷积层2: 32→32通道
            nn.Conv2d(32, 32, 5, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(2),
            # 卷积层3: 32→64通道
            nn.Conv2d(32, 64, 5, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(2),
            # 展平后接全连接层
            nn.Flatten(),
            nn.Linear(64*4*4, 64),
            nn.Linear(64, 10)
        )
    
    def forward(self, x):
        return self.net(x)

几个踩坑经验：

1. padding=2是为了保持特征图尺寸不变，配合kernel_size=5使用
2. 每个卷积层后要加激活函数，ReLU比Sigmoid训练更快
3. MaxPooling的步长默认为kernel_size，所以nn.MaxPool2d(2)等价于nn.MaxPool2d(2,2)

3. 训练过程优化

第一次训练时准确率卡在50%上不去，后来调整了三个关键点：

损失函数选择：CrossEntropyLoss自带Softmax，比手动写Softmax+NLLLoss更稳定

python复制loss_fn = nn.CrossEntropyLoss().cuda()  # 记得放到GPU上

优化器配置：Adam比SGD收敛更快，但最终准确率略低

python复制# SGD优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# 或者使用Adam
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

训练循环细节：

python复制for epoch in range(30):
    model.train()  # 训练模式
    for batch, (X, y) in enumerate(train_loader):
        X, y = X.cuda(), y.cuda()
        
        # 前向传播
        pred = model(X)
        loss = loss_fn(pred, y)
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # 每100批记录一次
        if batch % 100 == 0:
            print(f"Epoch:{epoch} | Batch:{batch} | Loss:{loss.item():.4f}")

用TensorBoard可视化训练过程特别实用：

python复制from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter()
writer.add_scalar('Loss/train', loss.item(), global_step=epoch*len(train_loader)+batch)

4. 测试与验证技巧

测试集评估要注意两个模式切换：

python复制model.eval()  # 切换评估模式
with torch.no_grad():  # 关闭梯度计算
    for X, y in test_loader:
        X, y = X.cuda(), y.cuda()
        pred = model(X)
        # 计算准确率...

验证单张图片时有三个预处理步骤容易出错：

1. 通道转换：PNG可能是RGBA四通道，要转RGB

python复制image = Image.open('plane.png').convert('RGB')

2. 尺寸调整：必须与训练尺寸一致

python复制transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor()
])

3. 维度扩展：模型需要4D输入(batch,channel,height,width)

python复制image = transform(image).unsqueeze(0).cuda()

保存模型时推荐同时保存结构和参数：

python复制# 保存完整模型
torch.save(model, 'full_model.pth')

# 只保存参数（需配合模型类使用）
torch.save(model.state_dict(), 'params_only.pth')

有个实际项目中遇到的坑：当验证准确率突然下降时，可能是数据没有shuffle导致的批次偏差。解决方法是在DataLoader中设置shuffle=True：

python复制test_loader = DataLoader(test_data, batch_size=64, shuffle=True)

我在笔记本上跑完30轮大约需要15分钟，最终测试准确率约65%。想要进一步提升可以尝试：

• 增加数据增强（旋转、裁剪等）
• 使用更复杂的网络结构（如ResNet）
• 调整学习率衰减策略
• 延长训练轮数到100轮以上