Pytorch神经网络搭建与训练：注意事项与避坑指南

初学可以参照这个项目，一步一步读懂代码：mnist-cnn-tutorial: 基于 PyTorch 的 MNIST 手写数字识别教学项目

本项目以 MNIST 手写数字识别为例，在实际开发过程中，从模型设计到训练调优，有许多关键细节决定了模型最终的性能。以下是神经网络开发中的核心注意事项。

1. 模型架构设计 (Architecture Design)

🟢 维度匹配 (Dimension Matching)

这是初学者最容易报错的地方。在 CNN 中，数据经过卷积和池化后，形状会发生变化。

• 注意点：进入全连接层（nn.Linear）前的特征数量必须手动计算正确。
• 代码参考：

  # 28x28 输入 -> 卷积1 -> 池化1 (14x14) -> 卷积2 -> 池化2 (7x7)
  # 最终通道数为 64，所以全连接层输入是 64 * 7 * 7 = 3136
  self.fc1 = nn.Linear(64 * 7 * 7, 128)
  self.relu3 = nn.ReLU()
  self.fc2 = nn.Linear(128, 10)
  

🟢 激活函数的选择

• 注意点：隐藏层通常使用 ReLU，它能缓解梯度消失问题。
• 避坑：最后一层输出层如果是多分类，通常不加激活函数（因为 PyTorch 的 CrossEntropyLoss 内部集成了 Softmax）。

2. 数据预处理 (Data Preprocessing)

🟢 归一化 (Normalization)

• 原理：将像素值从 [0, 255] 缩放到均值为 0、方差为 1 的分布。这能显著加快收敛速度。
• 代码参考：

  transform = transforms.Compose([
      transforms.ToTensor(),
      transforms.Normalize((0.1307,), (0.3081,)) # MNIST 数据集的官方均值和标准差
  ])
  

• 关键：推理时（GUI）必须使用与训练时完全相同的均值和标准差。

3. 训练循环 (Training Loop)

🟢 梯度清零 (Zero Grad)

• 注意点：PyTorch 的梯度是累加的。
• 必做操作：在每个 Batch 开始前必须调用 optimizer.zero_grad()。
• 代码参考：

  optimizer.zero_grad() # 必须清零，否则梯度会不断累加导致模型崩溃
  output = model(data)
  loss = criterion(output, target)
  loss.backward()
  optimizer.step()
  

🟢 训练模式与评估模式切换

• 注意点：某些层（如 Dropout, Batch Normalization）在训练和测试时的行为不同。
• 必做操作：
  • 训练时调用 model.train()
  • 验证/测试时调用 model.eval() 配合 with torch.no_grad(): 以节省显存和计算资源。

4. 损失函数与优化器 (Loss & Optimizer)

🟢 学习率调度 (Learning Rate Scheduler)

• 建议：训练后期降低学习率有助于模型进入更深的局部最小值。
• 代码参考：

  # 使用 StepLR，每 5 个 Epoch 将学习率衰减为原来的 0.1 倍
  optimizer = optim.Adam(model.parameters(), lr=0.001)
  scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
  

5. 推理与部署 (Inference)

🟢 输入维度的“伪造”

• 注意点：模型期望 [Batch, Channel, H, W] 4D 输入。
• 必做操作：单张图片推理时使用 unsqueeze(0) 增加 Batch 维度。
• 代码参考：

  # 图像缩放后是 (1, 28, 28)，通过 unsqueeze(0) 变成 (1, 1, 28, 28)
  input_tensor = transform(img_resized).unsqueeze(0).to(device)
  

🟢 置信度解析 (Softmax)

• 原理：Logits 输出不直观，需通过 torch.softmax 转化为概率。
• 注意：Softmax 后的最大值即为置信度。

  prob = torch.nn.functional.softmax(output, dim=1)
  confidence, predicted = torch.max(prob, 1)
  

总结：避坑清单

1. 检查设备：确保模型和数据都在同一个设备上（同为 cuda 或同为 cpu）。
2. Shuffle：训练集一定要 shuffle=True，验证集不需要。
3. Overfitting：如果训练 Loss 持续下降但验证准确率不动，考虑加入正则化或简化模型。