PyTorch 修改权重/字典 Key

sdlkjaljafdg

已于 2023-03-02 06:48:55 修改

阅读量814

点赞数

分类专栏： Pytorch 文章标签： pytorch 深度学习人工智能

于 2023-03-02 06:47:18 首次发布

本文链接：https://blog.csdn.net/Haku_yyf/article/details/129293251

版权

Pytorch 专栏收录该内容

31 篇文章 10 订阅

订阅专栏

Pytorch修改权重

torch.load

关于result = unpickler.load()ModuleNotFoundError: No module named ‘models‘解决

pytorch模型的加载与保存

import torch
from collections import OrderedDict

def mmlabmodel():
    mmlab_model = torch.load("epoch_30.pth", map_location=torch.device('cpu'))
    return mmlab_model
    # for key in model['state_dict']:
    #     print(model['state_dict'][key].shape)

def yolov5obb():
    import sys
    sys.path.insert(0, 'yolov5_obb/')

    """
    /home/hpc/PycharmProjects/Siam-NestedUNet/是一个项目路径，
    这个项目保存了权重checkpoint_epoch_59.pt,我把这个权重复制到另一个项目中加载
    需要加上注释上面2行的代码
    """
    ckpt = torch.load(r"best.pt", map_location=torch.device('cpu'))
    return ckpt
    # model = ckpt['model']
    # for key in model.state_dict():
    #     print(model.state_dict()[key].shape)

def mmlab2obb():
    mmlabmodels = mmlabmodel()
    obbmodels = yolov5obb()
    mmlabkeys = list(mmlabmodels['state_dict'].keys())
    obbkeys = list(obbmodels['model'].state_dict().keys())
    new_dict = OrderedDict()
    assert len(mmlabkeys) == len(obbkeys)
    for i in range(len(mmlabkeys)):
        new_dict[obbkeys[i]] = mmlabmodels['state_dict'][mmlabkeys[i]]
    obbmodels['model'].load_state_dict(new_dict)
    torch.save(obbmodels, 'obbnewsave.pt')

def obb2mmlab():
    mmlabmodels = mmlabmodel()
    obbmodels = yolov5obb()
    mmlabkeys = list(mmlabmodels['state_dict'].keys())
    obbkeys = list(obbmodels['model'].state_dict().keys())
    new_dict = OrderedDict()
    assert len(mmlabkeys) == len(obbkeys)
    for i in range(len(mmlabkeys)):
        new_dict[mmlabkeys[i]] = obbmodels['model'].state_dict()[obbkeys[i]]
    mmlabmodels['state_dict'] = new_dict
    torch.save(mmlabmodels, 'mmlabnewsave.pth')


if __name__ == '__main__':
    # mmlab2obb()
    obb2mmlab()
    print("Successful!!")