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

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

朽一

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

阅读量2.5k

点赞数

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

本文链接：https://blog.csdn.net/qq_43027065/article/details/121222475

版权

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

21 篇文章 15 订阅

订阅专栏

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")