PyTorch推理阶段冻结特征层模型的正确加载方式咨询

• 你当前的加载代码存在问题：model.load_state_dict(bestmodel['classifier'])这行是错误的。load_state_dict()需要接收与整个模型结构匹配的状态字典，而bestmodel['classifier']只是分类器层的参数，直接传入会导致参数key不匹配，触发报错。正确的做法是把分类器的参数加载到model.classifier上，而非整个模型。

• 推理时绝对不能设置pretrained=False：你训练时冻结了特征提取层，用的是预训练权重，推理时必须保持特征层权重和训练时一致。如果设pretrained=False，特征层会被随机初始化，和训练时的特征提取逻辑完全不同，推理结果必然出错。

• 正确的推理加载流程示例：

model = models.densenet161(pretrained=True)
for param in model.parameters():
    param.requires_grad = False
num_ftrs = model.classifier.in_features
model.classifier = torch.nn.Linear(num_ftrs, 2)
model.to(device)

# 加载最优模型参数
bestmodel_path = get_best_model(best)
bestmodel = torch.load(bestmodel_path)
# 将分类器参数加载到对应层
model.classifier.load_state_dict(bestmodel['classifier'])

# 切换到评估模式
model.eval()
with torch.no_grad():
    # 你的推理代码
    pass

补充：如果训练时你保存的是整个模型的状态字典（比如torch.save(model.state_dict(), save_path)），推理时可以更简化：初始化模型后直接调用model.load_state_dict(torch.load(bestmodel_path))即可，无需单独处理分类器。

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