model_list = [
                  './xxx.pt.tar',
                  './yyy.pt.tar',
                  ]
# 融合两个模型,模型一的backbone + 模型二的全连接层
def integration2(model_list, fl_model):

    worker_state_dict=[torch.load(x, map_location='cpu') if x.endswith('pt') else torch.load(x, map_location='cpu')['state_dict'] for x in model_list]
    print(worker_state_dict[0].keys())
    weight_keys=list(worker_state_dict[0].keys()) # ['features.0.weight', 'features.1.weight', 'features.1.bias'.....'output.1.weight', 'output.1.bias']
    fed_state_dict=collections.OrderedDict()
    for key in weight_keys:
        print('key is {}'.format(key))
        key_sum = 0
        for i in range(len(model_list)):
            key_sum += worker_state_dict[i][key]
        fed_state_dict[key] = key_sum /float(len(model_list))
    fl_model.load_state_dict(fed_state_dict)  # 融合后的模型
    return fl_model
Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐