Error initializing fasterrcnn_resnet50_fpn with num_classes and pretrained

[dipam0111]

🐛 Bug

In torchvision.models.detection.fasterrcnn_resent50_fpn there are two parameters, num_classes and pretrained.

If I pass just one of them, I am able to initialize the model

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(num_classes=6)

or

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

^this obv won't work for me since I have just 5 classes and not 91.

hence, when I pass both, I get this error

code:

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

error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/var/folders/5q/5dh5_83x359fqcb949hhtj8m0000gp/T/ipykernel_11119/3787706091.py in <module>
----> 1 model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained = True, num_classes=6)
      2 # num_classes = 6 # should be initialized as target_col.nunique + 1
      3 # in_features = model.roi_heads.box_predictor.cls_score.in_features
      4 # model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    370         state_dict = load_state_dict_from_url(model_urls['fasterrcnn_resnet50_fpn_coco'],
    371                                               progress=progress)
--> 372         model.load_state_dict(state_dict)
    373         overwrite_eps(model, 0.0)
    374     return model

   1404 
   1405         if len(error_msgs) > 0:
-> 1406             raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
   1407                                self.__class__.__name__, "\n\t".join(error_msgs)))
   1408         return _IncompatibleKeys(missing_keys, unexpected_keys)

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([91, 1024]) from checkpoint, the shape in current model is torch.Size([6, 1024]).
	size mismatch for roi_heads.box_predictor.cls_score.bias: copying a param with shape torch.Size([91]) from checkpoint, the shape in current model is torch.Size([6]).
	size mismatch for roi_heads.box_predictor.bbox_pred.weight: copying a param with shape torch.Size([364, 1024]) from checkpoint, the shape in current model is torch.Size([24, 1024]).
	size mismatch for roi_heads.box_predictor.bbox_pred.bias: copying a param with shape torch.Size([364]) from checkpoint, the shape in current model is torch.Size([24]).

To Reproduce

Steps to reproduce the behavior:

import torchvision

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

Expected behavior

It should not throw an error.

Environment

• PyTorch / torchvision Version (e.g., 1.0 / 0.4.0):

torch                             1.9.0
torchvision                       0.10.0

• OS (e.g., Linux):

MacOS

• How you installed PyTorch / torchvision (conda, pip, source):

pip

• Build command you used (if compiling from source):

None

• Python version:

python3 --version
>Python 3.8.5

• CUDA/cuDNN version:
• GPU models and configuration:
• Any other relevant information:

[NicolasHug]

I think that the pretrained model is pretrained for a specific number of classes, which is why we cannot set both. But it would be nice to have a better error message.

[dipam0111]

yes or you can change the final if statement to if pretrained and num_classes = 91: only then load the state dict

[datumbox]

All pre-trained models in TorchVision have a fixed number of classes that depends on the dataset that we used to train the model. If you want to define your own number of classes, you could load the pretrained model, do model surgery to overwrite those layers that make the class prediction and then fine-tune the model on your own data.

I'm closing the ticket as I believe addresses the question but feel free to reopen if you have more concerns.

[NicolasHug]

I had kept it open so as to not forget about the fact that we should raise a better error message, as mentioned above. It's fine to close though, I'll follow up on that soon(ish)

[datumbox]

It's definitely worth addressing it on the future but note it might be tricky to cover all corner-cases/parameters and return the right error message for each.

There are models that offer quite generic configuration and in theory someone can pass incompatible config for a specific pre-trained model. This is likely to be addressed more naturally when a proper versioning mechanism is introduced. There are ideas to link specific model configs to specific weights/versions but this is still under investigation.

[NicolasHug]

I agree it might be tricky to fully validate the kwargs, but for the non-kwargs argument it should be fairly straightforward, even if a bit verbose