Faster R-CNN模型推理阶段加载训练后权重的正确方式及性能差异疑问

问题背景 #

我想了解在推理阶段加载已训练模型保存权重的正确方式。

训练阶段的模型初始化与保存 #

训练时我使用预训练COCO权重和预训练ImageNet权重初始化模型：

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True, pretrained_backbone=True)

num_classes = 2 

# 获取分类器的输入特征数
in_features = model.roi_heads.box_predictor.cls_score.in_features

# 替换预训练的头部为新的预测器
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

model.to(device)

训练完成后，我将模型状态字典保存为：

'state_dict': model.state_dict()

推理阶段的权重加载尝试 #

训练完成后，我最初按如下方式加载权重：

# 初始化测试模型
model_test = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True, pretrained_backbone=True)
num_classes = 2
# 获取分类器的输入特征数
in_features = model_test.roi_heads.box_predictor.cls_score.in_features
# 替换预训练的头部为新的预测器
model_test.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
model_test.to(device)
    
bestmodel = get_best_model(args.best)
bestmodel = torch.load(bestmodel)
model_test.load_state_dict(bestmodel['state_dict'])

我原本认为自己训练的权重会覆盖初始预训练权重，但当我把测试模型初始化改为：

model_test = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone= False)

再加载最优模型权重时，性能会略有下降。

请问正确的权重加载方式是什么？如果理论上不应有差异，为何会出现这种情况？

解答 #

正确的权重加载方式 #

正确的加载逻辑是：先构建与训练阶段完全一致的模型结构，再直接加载训练好的state_dict，不需要先加载预训练权重。具体步骤如下：

1. 构建模型结构：和训练时一样，先初始化基础模型（pretrained参数设为True或False均可，只要结构完全匹配训练时的模型），然后替换头部的FastRCNNPredictor到对应类别数。
2. 加载训练保存的state_dict：用model.load_state_dict()加载自己训练好的权重，这会覆盖模型所有层的参数，包括backbone和头部。

标准的推理阶段代码示例：

# 构建和训练时完全一致的模型结构
model_test = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)
num_classes = 2
in_features = model_test.roi_heads.box_predictor.cls_score.in_features
model_test.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
model_test.to(device)

# 加载训练好的权重
bestmodel_path = get_best_model(args.best)
bestmodel = torch.load(bestmodel_path, map_location=device)
model_test.load_state_dict(bestmodel['state_dict'])
# 务必切换到评估模式
model_test.eval()

性能差异的可能原因 #

理论上两种初始化方式加载权重后，模型参数应该完全一致，性能差异大概率来自以下细节问题：

• 未切换到评估模式：如果推理时没有调用model_test.eval()，模型会保留训练模式的行为（比如BatchNorm的running均值/方差会动态更新、Dropout会生效），导致推理结果不稳定或性能下降。
• 设备不匹配：保存权重时模型在GPU，加载时直接在CPU（或反之），可能出现参数加载异常。建议用torch.load(bestmodel_path, map_location=device)确保权重加载到目标设备。
• 随机因素干扰：比如推理时未关闭训练式的数据增强、测试时未固定随机种子，或者BatchNorm的统计量在初始化时的差异，都可能导致性能波动。
• 权重保存/加载遗漏：比如保存时未完整保存state_dict（比如漏掉了BatchNorm的running_mean/running_var），或者加载时存在未匹配的参数键（可以用model.load_state_dict(..., strict=False)检查，但不建议长期使用）。

验证方法 #

可以提取两种初始化方式加载权重后，模型某一层的参数（比如backbone第一层卷积的权重）进行数值对比，如果参数完全一致，说明权重加载逻辑是正确的，性能差异就来自上述其他因素。

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