融合两个模型的权重 - 权重平均

最新推荐文章于 2024-04-07 17:52:53 发布

梦坠凡尘(AICV与前沿)

最新推荐文章于 2024-04-07 17:52:53 发布

阅读量3.5k

点赞数

分类专栏： Pytorch 文章标签： pytorch 深度学习机器学习

本文链接：https://blog.csdn.net/c2250645962/article/details/120023246

版权

Pytorch 专栏收录该内容

16 篇文章 33 订阅

订阅专栏

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