PyTorch基于ResNet50训练分类器触发输入维度不匹配RuntimeError

错误根因 #

ResNet50的网络结构在进入最后的fc全连接层之前，会经过全局平均池化层，输出维度为**[batch_size, 2048]的二维张量。你替换fc层时用了nn.Conv2d作为首层，而2D卷积要求输入是[batch_size, channel, height, width]**的四维张量，维度不匹配直接触发报错。
同时你的代码还有两处隐藏问题：

• ResNet50池化后的输出通道数为2048，你自定义卷积层的输入通道设为512，就算维度匹配也会触发通道数不匹配的报错
• 你做的是7分类单标签任务，用nn.BCELoss()搭配one-hot标签的逻辑不合理，且最终输出层没有加对应的激活函数，会导致loss计算异常

解决方案 #

如果你只是要做常规的ResNet迁移学习分类，直接把自定义头部的卷积层去掉，改用全连接层即可，符合你冻结预训练参数只训新增层的需求，修改后的核心代码如下：

# 替换你原有定义model.fc的部分
model = models.resnet50(pretrained=True)
for param in model.parameters():
    param.requires_grad = False
# 直接用全连接层作为分类头，不需要卷积层
model.fc = nn.Sequential(
    nn.Linear(2048, 512),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(512, 64),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(64, 7)
)
model.to(device)
# 损失函数改用交叉熵，不需要自己做one-hot编码，PyTorch会自动处理
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.fc.parameters(), lr=learning_rate)

同时要把训练循环里的one-hot编码代码删掉，CrossEntropyLoss直接接受原始标签即可：

# 删掉这行
# labels = torch.nn.functional.one_hot(labels, num_classes=7).float()

如果你确实需要在头部加卷积层，需要修改ResNet的整体结构，去掉最后的全局平均池化层，将最后一个残差块输出的四维特征直接传入自定义的卷积头，调整后代码如下：

model = models.resnet50(pretrained=True)
for param in model.parameters():
    param.requires_grad = False
# 去掉最后的平均池化和fc层，拿到前面的特征提取部分
backbone = nn.Sequential(*list(model.children())[:-2])
# 自定义头部，输入通道对应ResNet50最后输出的2048通道
head = nn.Sequential(
    nn.Conv2d(2048, 64, (3, 3), padding=1),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Conv2d(64, 64, (3, 3), padding=1),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Dropout(0.5),
    nn.Flatten(),
    # 这里的线性层输入维度要根据你输入图像尺寸计算，200x200输入的话经过ResNet和两次池化后对应维度是64*12*12=9216
    nn.Linear(9216, 512),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(512, 64),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(64, 7)
)
# 拼接成完整模型
model = nn.Sequential(backbone, head)
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(head.parameters(), lr=learning_rate)

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