Torchvision loads the weight of RESNET except the full connection layer

  Simply post how to load torchvision Pre training weight , Or always forget .

import torch
import torch.nn as nn
import torchvision

class ResNet(nn.Module):
    def __init__(self):
        super(ResNet, self).__init__()
        pass
    #  Go to ResNet Add weight inside 
    def init_weights(self, pretrained = True):
        """ Args: self:  Model itself  pretrained (bool) """
        if pretrained == True:
            #  obtain ResNet34 Pre training weight of 
            resnet34 = torchvision.models.resnet34(pretrained=True)
            pretrained_dict = resnet34.state_dict()
            """ load torchvision The pre training model and parameters in are passed after state_dict() Method extract parameters   Or directly from the official model_zoo download ： pretrained_dict = model_zoo.load_url(model_urls['resnet152'])"""
            #  Get the parameter Dictionary of the current model 
            model_dict = self.state_dict()
            #  take pretrained_dict It doesn't belong to model_dict Key out of 
            pretrained_dict = {
    k: v for k, v in pretrained_dict.items() if k in model_dict}
            #  Update existing model_dict
            model_dict.update(pretrained_dict)
            #  Load what we really need state_dict
            self.load_state_dict(model_dict)
            print(' Successfully loaded pre training weights ')
        else:
            pass
    
if __name__ == '__main__':
    resnet = ResNet()
    resnet.init_weights(pretrained=True)