pytorch 模型加权

pytorch 模型加权

from collections import OrderedDict

import torch

from LPRNetN.model.STN import STNet

STN = STNet()

state_dict1 = torch.load('weights/STNNet.pth', map_location=lambda storage, loc: storage)
state_dict2 = torch.load('weights/STNNet.pth', map_location=lambda storage, loc: storage)
new_state_dict1 = OrderedDict()
new_state_dict2 = OrderedDict()
for k, v in state_dict2.items():
    name = k.replace('module.', '')  # remove `module.`
    new_state_dict2[name] = v

for k, v in state_dict1.items():
    name = k.replace('module.', '')  # remove `module.`
    new_state_dict1[name] =0.5*new_state_dict2[name]+0.5* v

STN.load_state_dict(new_state_dict1)

打印模型参数：

下面两种结果是一样的：

if __name__ == '__main__':
    fc = nn.Linear(3, 5)
    params = list(fc.named_parameters())
    print(params.__len__())
    print(params[0])
    print(params[1])

    state_dict2=fc.state_dict()
    for k, v in state_dict2.items():
        print(k,v)