将gpu并行模型转换为单cpu/gpu模型

最新推荐文章于 2024-04-21 00:21:26 发布

朽一

最新推荐文章于 2024-04-21 00:21:26 发布

阅读量2.6k

点赞数

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

本文链接：https://blog.youkuaiyun.com/qq_43027065/article/details/121222475

版权

机器学习or深度学习专栏收录该内容

22 篇文章

订阅专栏

该代码实现了一个模型转换功能，通过加载指定的参数模型（para_model），将模型权重从模块化的形式转换为单个模型（single_model）。首先使用torch.load加载检查点，然后遍历检查点中的模型权重，移除'module.'前缀并保存到新的模型文件中。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import torch


def convert_model(para_model,single_model):
    checkpoint = torch.load(para_model ,map_location=torch.device("cpu"))

    output={}
    for key,value in checkpoint['model'].items():
        output[key.lstrip("module.")]=value

    torch.save({'model':output},single_model)

if __name__=="__main__":
    convert_model(para_model="model_2.pth", single_model="model.pth")