尝试加载 PyTorch 模型时出现大小不匹配运行时错误

Question

下面是我尝试运行的代码。fasterrcnn_foodtracker.pth是我正在尝试使用 PyTorch 加载的已经训练好的模型。

import torch
import torchvision
import cv2

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

path = '/home/amir/PycharmProjects/Food-Recognition/fasterrcnn_foodtracker.pth'
model.load_state_dict(torch.load(path, map_location=torch.device('cpu')), strict=False)
model.eval()

img = cv2.imread('twodishes.jpg')
prediction = model([img])
print(prediction)

出现大小不匹配的运行时错误。

RuntimeError: Error(s) in loading state_dict for FasterRCNN:
    size mismatch for roi_heads.box_predictor.cls_score.weight: copying a param with shape torch.Size([100, 1024]) from checkpoint, the shape in current model is torch.Size([91, 1024]).
    size mismatch for roi_heads.box_predictor.cls_score.bias: copying a param with shape torch.Size([100]) from checkpoint, the shape in current model is torch.Size([91]).
    size mismatch for roi_heads.box_predictor.bbox_pred.weight: copying a param with shape torch.Size([400, 1024]) from checkpoint, the shape in current model is torch.Size([364, 1024]).
    size mismatch for roi_heads.box_predictor.bbox_pred.bias: copying a param with shape torch.Size([400]) from checkpoint, the shape in current model is torch.Size([364]).

Answer 1

在我看来，您的模型配置与模型检查点的内容不匹配。我想您的模型具有诸如input_size、output_size（或num_classes）之类的参数，这些参数最终将定义您的模型的外观（每层的参数数量等）。

快速示例：

# two layer linear 'network', `num_in` input units, `num_out` output units
model = nn.Sequential(nn.Linear(num_in, 100), nn.Linear(100, num_out))

num_in = 10如果您为和训练并保存此模型num_out = 20，将这些参数更改为num_in = 12/num_out = 22并加载之前保存的模型，加载例程将抱怨形状不匹配（10 与 12 和 20 与 22）。

这似乎就是发生在你身上的事情。解决方案：您需要确保使用与用于生成正在加载的检查点的模型相同的超参数来初始化模型。