PyTorch MNIST数字分类模型加载后预测结果恒为5问题求助

1. 模型保存与加载流程校验 #

• 确认保存/加载方式匹配：训练时若用torch.save(model.state_dict(), 'model.pth')保存参数，加载时必须先实例化与训练时结构完全一致的模型，再执行model.load_state_dict(torch.load('model.pth'))；若直接保存整个模型（torch.save(model, 'model.pth')），加载时需确保模型类的定义在当前环境中可被正确识别。
• 强制切换到评估模式：加载模型后必须调用model.eval()，关闭训练时启用的Dropout、BatchNorm等层的动态行为，否则这些层的随机特性会导致输出异常。
• 核对模型结构：重新搭建模型时，需严格对齐训练时的所有细节——包括层数、激活函数、Dropout率、BatchNorm的affine参数等，哪怕细微差异都会导致参数加载错位，引发输出异常。

2. 输入数据预处理一致性检查 #

• 匹配训练时的归一化逻辑：MNIST训练时通常会做像素值归一化（如transforms.Normalize((0.1307,), (0.3081,))），预测时必须对输入执行完全相同的变换，否则输入数据分布偏离训练集，会导致模型输出失效。
• 确保输入维度正确：PyTorch模型期望输入维度为(batch_size, channels, height, width)，单张图片需调整为(1, 1, 28, 28)而非原始的(28,28)，维度错误会导致模型计算逻辑混乱。
• 对齐数据类型：模型参数默认是float32，输入数据需转换为同类型（如input = input.float()），避免因类型不匹配导致的计算异常。

3. 推理环节细节排查 #

• 统一设备环境：若训练时使用GPU，加载模型后需执行model.to(device)，同时将输入数据也移至同一设备（input = input.to(device)），跨设备计算会导致输出异常。
• 查看原始输出Logits：不要只看argmax的结果，打印模型输出的原始Logits值。如果所有样本的Logits中对应数字5的权重远高于其他类别，说明模型参数加载后存在问题；若Logits分布正常但argmax结果错误，则需检查后处理代码逻辑。
• 关闭梯度计算：推理时用with torch.no_grad():包裹预测代码，虽然这一般不会导致输出一致，但能避免不必要的内存占用和潜在的计算干扰。

4. 极端场景验证 #

• 检查模型文件完整性：重新保存一次训练好的模型，加载后对比训练结束时与加载后的模型参数（如打印model.fc.weight[:1]），确认参数未损坏或错位。
• 验证测试数据有效性：随机选取几张测试图片可视化，确认输入确实是不同数字的样本，排除测试数据被错误处理为同一类的可能。

内容的提问来源于stack exchange，提问作者grSage