PyTorch中的torchscript模型转换为普通的nn.Module模型

john_bh

于 2024-06-03 08:38:41 发布

阅读量245

点赞数 1

分类专栏： PyTorch 文章标签： pytorch 人工智能 python

本文链接：https://blog.youkuaiyun.com/john_bh/article/details/139401951

版权

PyTorch 专栏收录该内容

21 篇文章

订阅专栏

PyTorch中的torchscript模型转换为普通的nn.Module模型

from models import *
def ptl2pt():
    map_location = 'cpu' #如果要在GPU上运行，则此处设置为 'cuda'
    model = torch.jit.load("test.ptl",map_location = map_location)

    new_model = Net(num_feats=32, depth_chanels=1, color_channel=3, kernel_size=3).to(device)
    state_dict = model.state_dict() #从torchscript 模型中提取参数
    new_model.load_state_dict(state_dict) # 将参数加载到新模型中
    torch.save(new_model.state_dict(),"test.pth")# 保存新模型