PyTorch加载已保存CapsuleNet模型输出全为0问题排查

我之前在调试CapsNet的时候也碰到过类似的输出全0问题，结合你的操作步骤，大概率是这几个环节出了问题，给你逐一拆解排查：

1. 模型配置必须完全对齐 #

你训练时用的Config()参数，和加载时实例化CapsNet(Config())的配置必须丝毫不差——比如胶囊的维度、卷积层的通道数、分类数这些核心参数，只要有一个不一样，模型结构就不匹配，加载完state_dict后参数就会乱掉，自然输出异常。

• 解决办法：把你训练时用的Config代码原封不动复制到加载脚本里，别改任何参数，哪怕是看起来无关的小配置都不行。

2. CPU/GPU设备不匹配是重灾区 #

如果训练时用GPU跑的，保存的模型参数默认是存在GPU设备上的，换到另一台机器如果没GPU，或者没指定加载位置，就会导致参数加载异常；反过来训练用CPU，加载用GPU也可能出问题。

• 解决办法：加载时明确指定设备：

  # 要是当前机器只有CPU
  capsnet.load_state_dict(torch.load('capsnet_mnist_state.pt', map_location=torch.device('cpu')))
  # 有GPU的话就指定cuda
  capsnet.load_state_dict(torch.load('capsnet_mnist_state.pt', map_location=torch.device('cuda')))
  

  另外别忘了把模型和输入数据移到同一个设备上：

  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  capsnet = capsnet.to(device)
  # 测试输入数据也要同步过去
  test_input = test_input.to(device)
  

3. 数据预处理要和训练时完全一致 #

很多人容易忽略这点：如果测试时的图像预处理和训练时不一样（比如归一化的均值/标准差、图像尺寸、通道顺序），模型拿到的输入分布不对，也会输出奇怪的结果。

• 解决办法：翻出你训练时的DataTransform代码，比如训练时用了：

  transform = transforms.Compose([
      transforms.ToTensor(),
      transforms.Normalize((0.1307,), (0.3081,))
  ])
  

  测试时必须用一模一样的transform，不能自己随便改。

4. 多GPU训练的state_dict有前缀坑 #

要是你训练时用了nn.DataParallel（多GPU并行训练），那保存的state_dict里的参数键会多一个module.前缀，直接加载到普通的CapsNet实例里就会匹配失败，模型还是初始随机参数，输出全0就很正常了。

• 解决办法：加载时去掉这个前缀：

  state_dict = torch.load('capsnet_mnist_state.pt')
  # 批量替换键名，去掉module.前缀
  cleaned_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
  capsnet.load_state_dict(cleaned_state_dict)
  

  或者先用strict=False检查哪些键不匹配，再针对性处理：

  capsnet.load_state_dict(torch.load('capsnet_mnist_state.pt'), strict=False)
  

  （不过strict=False要谨慎用，可能会掩盖真正的结构不匹配问题）

5. 别搞错CapsNet的输出解读 #

CapsNet的输出不是直接的类别概率，而是每个类别的胶囊向量，你得计算向量的L2范数才能得到对应类别的置信度——要是你直接输出原始的胶囊向量，或者没做这一步，可能会误以为输出全0。

• 解决办法：检查你的输出处理代码，应该是这样的：

  outputs = capsnet(test_input)
  # 计算每个胶囊的长度作为类别概率
  class_probs = torch.norm(outputs, dim=-1)
  # 取最大概率对应的类别
  predictions = class_probs.argmax(dim=1)
  

建议你按这个顺序排查，先查配置和设备，这俩是最常见的问题。

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