Error Importing torchvision.models.resnet34 with num_classes parameter

I saw that the resnet34 which inherits from ResNet class has a parameter called num_classes which makes the last fc layer have output units equal to num_classes.
I tried creating a model of resnet34

new_model = models.resnet34(pretrained=True,num_classes=14)

I got this error

RuntimeError                              Traceback (most recent call last)
~/.conda/envs/my_root/lib/python3.6/site-packages/torch/nn/modules/ in load_state_dict(self, state_dict, strict)
    481                 try:
--> 482                     own_state[name].copy_(param)
    483                 except Exception:

RuntimeError: inconsistent tensor size, expected tensor [14 x 512] and src [1000 x 512] to have the same number of elements, but got 7168 and 512000 elements respectively at /opt/conda/conda-bld/pytorch_1512386481460/work/torch/lib/TH/generic/THTensorCopy.c:86

During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)
<ipython-input-31-fd5331f75d50> in <module>()
----> 1 new_model = models.resnet34(pretrained=True,num_classes=14)
      2 # for i,param in enumerate(model_ft.parameters()):
      3 #     if i ==  10:
      4 #         break
      5 #     param.requires_grad = False

~/.conda/envs/my_root/lib/python3.6/site-packages/torchvision-0.2.0-py3.6.egg/torchvision/models/ in resnet34(pretrained, **kwargs)
    174     model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
    175     if pretrained:
--> 176         model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
    177     return model

~/.conda/envs/my_root/lib/python3.6/site-packages/torch/nn/modules/ in load_state_dict(self, state_dict, strict)
    485                                        'whose dimensions in the model are {} and '
    486                                        'whose dimensions in the checkpoint are {}.'
--> 487                                        .format(name, own_state[name].size(), param.size()))
    488             elif strict:
    489                 raise KeyError('unexpected key "{}" in state_dict'

RuntimeError: While copying the parameter named fc.weight, whose dimensions in the model are torch.Size([14, 512]) and whose dimensions in the checkpoint are torch.Size([1000, 512]).

I could have simply created the model and then change the last layer like this

new_model = models.resnet34(pretrained=True)
new_model.fc = nn.Linear(2048,14)

Any solution?

Your second approach looks good!
This is the way to use pretrained weights and a custom last layer.
It’s also used in the transfer learning tutorial.

1 Like