PyTorch CNN模型参数无法正常存储问题排查与解决问询

背景与问题描述 #

需求是构建支持查看特征图的图像分类CNN模型，原模型将卷积、池化等模块存储在普通字典self.architecture中，通过forward方法依次执行并保存各层输出到self.maps，用于后续查看特征图。

原模型定义代码：

class ConvNet(nn.Module):
  def __init__(self, input_channels, output_dim):
    super().__init__()
    # input 48x48
    self.architecture = {
        "conv1": self.convblock(input_channels, 128, (3,3)), # 46x46
        "conv2" : self.convblock(128, 64, (3, 3), bnorm=True), # 44x44
        "pool1" : self.poolblock((2,2)), # 22x22
        "conv3" : self.convblock(64, 16, (3,3), stride=2), #10x10
        "conv4" : self.convblock(16, 10, (3,3)), # 8x8
        "pool2" : self.poolblock((2,2), bnorm=10), # 4x4
        "feedforward" : nn.Sequential(
          nn.Flatten(), # 4x4x10 = 160
          nn.Linear(160, 128), # 128
          nn.ReLU(inplace=True),
          nn.Dropout(0.3),
          nn.Linear(128, output_dim), # 3
          nn.Softmax(dim=1)
        )                  
    }
    self.maps = {}

  def forward(self, x):
    image = x
    for name, layer in self.architecture.items():
      out = layer(image)
      self.maps[name] = out
      image = out
    return image
     
  def convblock(self, inp, out, kernel, stride=1, bnorm=False):
    if bnorm:
      return nn.Sequential(
        nn.Conv2d(inp, out, kernel, stride=stride),
        nn.ReLU(inplace=True),
        nn.BatchNorm2d(out)
      )
    else:
      return nn.Sequential(
        nn.Conv2d(inp, out, kernel, stride=stride),
        nn.ReLU(inplace=True)
      )

  def poolblock(self, kernel, bnorm=None):
    if bnorm is None:
      return nn.MaxPool2d(kernel)
    else:
      return nn.Sequential(
          nn.MaxPool2d(kernel),
          nn.BatchNorm2d(bnorm)
      )

  def get_map(self, im, layer):
    fig, ax = plt.subplots(1,2, figsize=(20,10), gridspec_kw={'width_ratios': [1,3]})
    ax[0].set_xticks([])
    ax[0].set_yticks([])
    ax[0].imshow(im.reshape(im.shape[-2],im.shape[-1],1), cmap="gray") # Shows Input image
    self(im)
    map = self.maps[layer]
    map=map.reshape(map.shape[1],1,map.shape[-2],map.shape[-1])
    ax[1].set_xticks([])
    ax[1].set_yticks([])
    rows = max(int(map.shape[0]/8), 8)
    ax[1].imshow(make_grid(map,nrow=rows).permute(1, 2, 0)) # Shows all the channels after an operation.

实例化模型与优化器时：

model = ConvNet(1, 3).to(device)
adam_opt = torch.optim.Adam(model.parameters(), lr=learning_rate)

出现错误：

/usr/local/lib/python3.10/dist-packages/torch/optim/optimizer.py in __init__(self, params, defaults)
    271         param_groups = list(params)
    272         if len(param_groups) == 0:
--> 273             raise ValueError("optimizer got an empty parameter list")
    274         if not isinstance(param_groups[0], dict):
    275             param_groups = [{'params': param_groups}]

ValueError: optimizer got an empty parameter list

问题根源 #

这种架构不符合PyTorch子类化nn.Module的参数追踪规则：

• PyTorch的nn.Module只会自动识别并注册直接作为实例属性的子模块（如self.conv1 = ...），这些子模块会被加入到模块内部的_modules字典中。
• 普通Python字典（如原代码中的self.architecture）不会被nn.Module处理，字典内的子模块不会被注册到_modules，因此调用model.parameters()时无法收集到任何参数，导致优化器初始化失败。

修改方案 #

最简洁的修改是用PyTorch提供的nn.ModuleDict替代普通字典，nn.ModuleDict是专门用于存储子模块的容器类，会自动将内部的子模块注册到_modules中，同时保留字典的迭代和访问特性。

修改__init__方法中的self.architecture定义：

self.architecture = nn.ModuleDict({
    "conv1": self.convblock(input_channels, 128, (3,3)), # 46x46
    "conv2" : self.convblock(128, 64, (3, 3), bnorm=True), # 44x44
    "pool1" : self.poolblock((2,2)), # 22x22
    "conv3" : self.convblock(64, 16, (3,3), stride=2), #10x10
    "conv4" : self.convblock(16, 10, (3,3)), # 8x8
    "pool2" : self.poolblock((2,2), bnorm=10), # 4x4
    "feedforward" : nn.Sequential(
      nn.Flatten(), # 4x4x10 = 160
      nn.Linear(160, 128), # 128
      nn.ReLU(inplace=True),
      nn.Dropout(0.3),
      nn.Linear(128, output_dim), # 3
      nn.Softmax(dim=1)
    )                  
})

修改后，forward方法无需改动（nn.ModuleDict的迭代方式和普通字典一致），model.parameters()就能正确收集所有子模块的参数，优化器初始化即可正常运行，同时原有的特征图保存和查看逻辑也不受影响。

另一种可选方案是将每个子模块单独注册为实例属性（如self.conv1 = ...、self.conv2 = ...），但这种方式会破坏原有的字典管理结构，不如nn.ModuleDict适配性好。

PyTorch存储类参数的机制 #

• 子模块注册：nn.Module内部维护一个_modules有序字典，所有直接赋值给实例的子模块（如self.conv1 = nn.Conv2d(...)）都会被自动添加到_modules中，包括nn.Sequential、nn.ModuleDict等容器类内部的子模块。
• 参数收集：当调用model.parameters()时，会递归遍历_modules中的所有子模块，收集所有类型为nn.Parameter的可训练参数，以及子模块中包含的参数。
• 容器类作用：nn.ModuleDict、nn.ModuleList、nn.Sequential等都是nn.Module的子类，它们的核心作用是让PyTorch能够识别并管理内部的子模块，确保参数被正确追踪和更新。

内容的提问来源于stack exchange，提问作者Ku-hello