pytorch迁移部分模型参数到新模型

最新推荐文章于 2024-04-08 10:37:00 发布

EastWind2017

最新推荐文章于 2024-04-08 10:37:00 发布

阅读量832

点赞数 1

分类专栏： PYTHON

本文链接：https://blog.csdn.net/h1534589653/article/details/127868755

版权

pytorch 深度学习人工智能

PYTHON 专栏收录该内容

3 篇文章 0 订阅

订阅专栏

def func_converse():
    heads = {'hm': 1, 'vaf': 2, 'haf': 1}
    model = deletedres.ResBiNet(n_classes=1, heads=heads)
    save_model = torch.load(r'./resbinet/ressgd_child.pth')
    model_dict = model.state_dict()
    state_dict = {k: v for k, v in save_model.items() if k in model_dict.keys()}
    print(state_dict.keys())  # dict_keys(['w', 'conv1.weight', 'conv1.bias', 'conv2.weight', 'conv2.bias'])
    model_dict.update(state_dict)
    model.load_state_dict(model_dict)
    torch.save(model.state_dict(), 'sgd.pth')