Size Mismatch Runtime Error When Trying to Load a PyTorch Model

Question

Below is the code I am trying to run. fasterrcnn_foodtracker.pth is the already trained model I am trying to load with 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)

A runtime error appears with a size mismatch.

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

As mentioned in previous answers, the error happens because the checkpoint contains weights for a different number of classes.

Setting strict=False when loading the state_dict does not fix the error because then only the layers with key mismatches are skipped during the loading.

You can find here a solution for the error, which involves first removing the layers with mismatched sizes from the state_dict and only then loading it with strict=False.

In your case, it would look like this:

current_model_dict = model.state_dict()
loaded_state_dict = torch.load(path, map_location=torch.device('cpu'))
new_state_dict={k:v if v.size()==current_model_dict[k].size()  else  current_model_dict[k] for k,v in zip(current_model_dict.keys(), loaded_state_dict.values())}
model.load_state_dict(new_state_dict, strict=False)