pytorch 忽略層權重

列印網絡層結構：

if "bias" not in name and "batch_norm" not in name:
      print(name.replace(" ","_"),str(list(v.size())))
           

忽略層權重：

from collections import OrderedDict

import torch

ignor_names=["ssh1.conv7x7_3.1",
             "ssh2.conv3X3.1"]

if __name__ == '__main__':

    trained_model=r"../weights/0.135.pth"
    state_dict = torch.load(trained_model)
    # create new OrderedDict that does not contain `module.`

    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        head = k[:7]
        if head == 'module.':
            name = k[7:]  # remove `module.`
        else:
            name = k
        is_pass=True
        for ignor_name in ignor_names:
            if ignor_name in name:
                is_pass=False
        if is_pass:
            print(k)
            new_state_dict[name] = v

    if n_gpu > 1:
        net.module.load_state_dict(new_state_dict, strict=False)
    else:
        net.load_state_dict(new_state_dict)