How to load part of pre trained model?

albertxavier001: 

chenyuntc: I’m afraid not <SCODE>    def load_my_state_dict(self, state_dict):
 
        own_state = self.state_dict()
        for name, param in state_dict.items():
            if name not in own_state:
                 continue
            if isinstance(param, Parameter):
                # backwards compatibility for serialized parameters
                param = param.data
            own_state[name].copy_(param)
<ECODE>

apaszke: You can remove all keys that don’t match your model from the state dict and use it to load the weights afterwards: <SCODE>pretrained_dict = ...
model_dict = model.state_dict()

# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict) 
# 3. load the new state dict
model.load_state_dict(pretrained_dict)
<ECODE>

albertxavier001: 

albertxavier001: <SCODE>model.load_state_dict(model_dict)<ECODE>

apaszke: Yes it should. I edited the snippet.

cham-3: Thanks. I am looking for it. It seems awesome!!!

ImgPrcSng: Consider the situation where I would like to restore all weights till the last layer. <SCODE># args has the model name, num classes and other irrelevant stuff
self._model = models.__dict__[args.arch](pretrained = False, 
                                         num_classes = args.classes, 
                                         aux_logits = False)


if self.args.pretrained:
              
    print("=> using pre-trained model '{}'".format(args.arch))
    pretrained_state = model_zoo.load_url(model_names[args.arch])
    model_state = self._model.state_dict()

    pretrained_state = { k:v for k,v in pretrained_state.iteritems() if k in model_state and v.size() == model_state[k].size() }
    model_state.update(pretrained_state)
    self._model.load_state_dict(model_state)
<ECODE> Shouldn’t we also be checking if the sizes match before restoring? It looks like we are comparing only the names.

smth: 

BrianHuang: model.load_state_dict(model_dict) step 3 should look like this

chenchr: 

karttikeya_mangalam: 

zeakey: Assum following situation: <SCODE>pretrained_dict: ['A', 'B', 'C', 'D']
model_dict: ['A', 'B', 'C', 'E']
<ECODE> <SCODE>pretrained_dict: ['A', 'B', 'C']
model_dict: ['A', 'B', 'C', 'E']
<ECODE> The complete snippet is therefore as follow: <SCODE>pretrained_dict = ...
model_dict = model.state_dict()

# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict) 
# 3. load the new state dict
model.load_state_dict(model_dict)
<ECODE>

romina.ui: 

ptrblck: 

sriharsha0806: Is it param = Parameter.data ?

deJQK: It might be a little late to ask. If I want to skip some layers, like if I train the model with batch normalization, but want to use the trained bn version for that without batch normalization, how can I change the layers’ names? Because otherwise, the name might be different, and it will complain about size mismatching.

kingxueyuf: This is very useful, thanks

nurlano: Shouldn’t the # 3  be <SCODE>model.load_state_dict(model_dict)
<ECODE> instead of <SCODE>model.load_state_dict(pretrained_dict)
<ECODE> ?

Zichun_Zhang: