如何修复PyTorch神经网络训练时出现的通道尺寸不匹配问题

错误1：unsqueeze(dim=3)触发通道数不匹配报错 #

原因 #

PyTorch的Conv2d层要求输入遵循(N, C, H, W)的格式，分别对应批次大小、通道数、特征高度、特征宽度。你在dim=3位置新增维度后，输入形状变为(128, 160, 101, 1)，卷积层读取第二维作为通道数，得到160通道，和你定义的Conv1要求的1输入通道不匹配。

解决 #

在dim=1位置新增通道维度，修改后输入形状为(128, 1, 160, 101)，符合Conv2d的输入要求。

错误2：unsqueeze(dim=1)后触发IndexError #

原因 #

你定义的LeNet模型forward方法直接返回预测张量，不是键值对字典，你写的(net(inputs))['out']是针对返回字典的模型（比如torchvision内置的分割模型）的取值方式，对普通张量做字典索引自然会报错。

解决 #

将推理代码修改为：

inputs = inputs.unsqueeze(1)
outputs = net(inputs)

删除多余的['out']索引操作。

错误3：反向传播阶段触发梯度形状不匹配报错 #

原因 #

你当前LeNet的全连接层fc1的输入维度是硬编码的16*5*5=400，这个数值是适配原始LeNet输入2828手写数字数据集的计算结果。而你的输入是1160*101的频谱数据，经过两次卷积、两次池化后的flatten维度为13024，和硬编码的400不匹配，因此前向传播时全连接层输入形状不匹配，反向传播时报梯度形状错误。
维度计算过程：

• 输入形状：(128, 1, 160, 101)
• 经过conv1+2*2池化后输出形状：(128, 6, 78, 48)
• 经过conv2+2*2池化后输出形状：(128, 16, 37, 22)
• flatten后维度：16 * 37 * 22 = 13024，和报错提示的期望维度一致。

解决 #

修改LeNet的fc1层输入维度，将：

self.fc1   = nn.Linear(16*5*5, 120)

替换为：

self.fc1   = nn.Linear(16*37*22, 120)

如果要避免硬编码维度导致后续修改输入大小出错，也可以用动态计算的方式写forward：

def forward(self, x):
    out = F.relu(self.conv1(x))
    out = F.max_pool2d(out, 2)
    out = F.relu(self.conv2(out))
    out = F.max_pool2d(out, 2)
    out = torch.flatten(out, 1) # 自动flatten除了批次维度之外的所有维度
    out = F.relu(self.fc1(out))
    out = F.relu(self.fc2(out))
    out = self.fc3(out)
    return out

额外注意点 #

你当前的训练循环中，没有给从dataloader拿到的inputs加通道维度，需要在mixup操作之前补充unsqueeze操作，修改训练循环对应部分：

if use_cuda:
    inputs, targets = inputs.cuda(), targets.cuda()
# 新增通道维度
inputs = inputs.unsqueeze(1)
inputs, targets_a, targets_b, lam,layer, cost = mixup_data(inputs, targets, args.alpha,args.mixupBatch, use_cuda)

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