使用预训练VGG16进行5分类任务时CrossEntropyLoss权重维度不匹配RuntimeError问题排查

嘿，我来帮你搞定这个错误！咱们一步步拆解问题：

错误根源 #

你遇到的RuntimeError本质是模型输出维度和损失函数的权重维度不匹配，问题出在两个细节上：

1. 误用VGG16的分类器属性：VGG系列模型的最后分类层属性是model.classifier，而不是model.fc（这个属性是ResNet等模型常用的）。你现在的代码只是给模型新增了一个fc属性，但原来的classifier层仍然输出1000维的logit（对应ImageNet的1000分类）。
2. 替换层的输入维度错误：就算你改对了属性名，你写的nn.Linear(1000,512)也不对——VGG16的特征提取部分输出的是4096维特征，这才是原classifier第一层的输入维度，不是1000维。

这两个问题叠加，导致模型实际输出还是1000维，但损失函数的class_weights是针对5类的，自然会触发“权重应该对应1000类或无类，却得到5类权重”的错误。

解决方法 #

1. 正确替换VGG16的分类器层 #

把model.fc改成model.classifier，同时修正第一层Linear的输入维度为4096：

model = models.vgg16(pretrained=True) # 下载在Imagenet数据集上预训练的vgg16模型
# 替换预训练vgg16的classifier层为5个新层（注意输入维度是4096）
model.classifier = nn.Sequential(
    nn.Linear(4096, 512), 
    nn.ReLU(inplace=True), 
    nn.Linear(512,256), 
    nn.ReLU(inplace=True), 
    nn.Linear(256,128), 
    nn.ReLU(inplace=True), 
    nn.Linear(128,64), 
    nn.ReLU(inplace=True), 
    nn.Linear(64,5)
)

2. （可选）冻结预训练特征层（迁移学习最佳实践） #

如果是做迁移学习，建议冻结特征提取部分的参数，避免破坏预训练的权重，只训练新替换的分类器：

# 冻结features层的所有参数
for param in model.features.parameters():
    param.requires_grad = False

3. 验证模型输出维度 #

可以加一行测试代码，确认模型输出是5维：

# 生成一个符合VGG16输入要求的随机张量（batch_size=1, 3通道, 224x224）
test_input = torch.randn(1, 3, 224, 224).to(device)
output = model(test_input)
print(output.shape) # 正常输出应为 torch.Size([1, 5])

额外提示 #

你的类别权重计算逻辑是没问题的，修正模型结构后，CrossEntropyLoss(weight=class_weights)就能正常工作了。

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