加载PyTorch训练的FasterRCNN模型时出现Size Mismatch运行时错误

错误根因 #

• 官方默认的fasterrcnn_resnet50_fpn模型基于COCO数据集训练，默认输出91个类别，对应检测头参数维度为分类层[91, 1024]、回归层[364, 1024]
• 你加载的fasterrcnn_foodtracker.pth是100类食品检测场景的训练权重，对应检测头参数维度为分类层[100, 1024]、回归层[400, 1024]，二者维度不匹配，因此触发报错

修复方案 #

修改模型初始化逻辑，替换检测头为和训练权重匹配的100类配置即可，完整可运行代码如下：

import torch
import torchvision
import cv2
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

# 初始化基础模型，无需保留原生COCO预训练权重可将pretrained设为False
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False)
# 提取检测头的输入特征维度
in_channels = model.roi_heads.box_predictor.cls_score.in_features
# 替换检测头为100类配置，和待加载权重维度对齐
model.roi_heads.box_predictor = FastRCNNPredictor(in_channels, num_classes=100)

# 加载自定义训练权重
weight_path = '/home/amir/PycharmProjects/Food-Recognition/fasterrcnn_foodtracker.pth'
model.load_state_dict(torch.load(weight_path, map_location=torch.device('cpu')), strict=False)
model.eval()

# 读取并预处理图片，适配模型输入要求
img = cv2.imread('twodishes.jpg')
# 转换为RGB通道（cv2默认BGR）
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# 转换为0-1范围的float张量，通道维度前置
input_tensor = torch.tensor(img_rgb / 255.0, dtype=torch.float).permute(2, 0, 1)

# 推理
prediction = model([input_tensor])
print(prediction)

补充说明 #

你原来的代码直接传入cv2读取的uint8格式BGR图片也能运行，但不符合模型训练时的输入规范，容易出现检测精度异常的问题，上面的代码新增了输入预处理步骤，可保证推理结果和训练时一致。

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