Size Mismatch Runtime Error When Trying to Load a PyTorch Model

user_9050888 0
Viewed 4383

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]).
pytorch python
2 Answers

It seems to me that your model configuration does not match the content of the model checkpoint. I imagine your model has parameters such as input_size, output_size (or num_classes) which will eventually define how your model looks like (number of parameters per layer, etc.).

Quick example:

# 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))

If you train and save this model for num_in = 10 and num_out = 20, change these parameters to, say, num_in = 12 / num_out = 22 and load your previously saved model, the load routine will complain that the shapes do not match (10 vs. 12 and 20 vs. 22).

This seems to be what is happening to you. The solution: You need to make sure to initialize your model with the same hyperparameters of the model used to generate the checkpoint you are loading.

user_16439390

Try and load as pretrained=False, then update the state dict.. This is because RCNN model was trained on x number of classes, but you are trying to match with weights from another model which was probably trained on different number of classes.

user_149984590
Related