How to use a pretrain model to update part of the model in torchvison

vision
1

Im using FCN50+Resnet50 from torchvision, but it does not have pretrain model available now.
Is there a way to use the pretraining weights from resnet50 to just update part of the FCN50 which belongs to Resnet?
Here is what I mean



model_1 = torchvision.models.segmentation.fcn_resnet50(pretrained=False, num_classes=21).to('cuda:0')
 resnet50 = models.resnet50(pretrained=True)

How can I just update part of model1 which is in common which model2?

2

Hi,
You just have to create new state_dict from resnet50 model which doesn’t contain last fully connected classification layer and load it to the fcn_resnet50.backbone.

import torchvision
from collections import OrderedDict
m1 = torchvision.models.segmentation.fcn_resnet50(pretrained=False, num_classes=21)
m2 = torchvision.models.resnet50(pretrained=True)

new_state_dict = OrderedDict()
for k, v in m2.state_dict().items():
    if k != "fc.weight" and k != "fc.bias":
        new_state_dict[k] = v

m1.backbone.load_state_dict(new_state_dict)
1 Like