record-----torch(.pth)模型转onnx(.onnx)模型的方法

最新推荐文章于 2024-03-28 17:27:43 发布
最新推荐文章于 2024-03-28 17:27:43 发布
阅读量1.3k 收藏 5
点赞数
分类专栏： 人工智能
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/ZXF_1991/article/details/115676795
版权

torch转onnx

前言

torch转onnx目的用于将torch(.pth)模型转为tensorRT(.engine)，pth->onnx->engine

code

"""
需要的环境：pytorch1.4

进行模型转换，torch->onnx
环境要求:tvm转torch模型必须在pytorch1.4的环境下进行，1.4下可以进行torch->onnx->tvm,
同时也可以进行torch->libtorch转换，但是c++中源码编译的libtorch库要和torch->libtorch转换时使用的pytorch版本一致
由于c++工程中使用的libtorch库是1.0,因此torch->libtorch的过程也要在pytorch1.0环境下进行
"""
import sys
#如果将编译的tvm添加到系统环境变量中不起作用，需要在此进行再次系统添加
# sys.path.insert(0,'/data_1/software/use_software/tvm/tvm/topi/python/')
# sys.path.insert(0,'/data_1/software/use_software/tvm/tvm/python/')
# sys.path.insert(0,'/data_1/software/use_software/tvm/tvm/nnvm/')
import torch
import torch.onnx  # 谨记：torch.onnx包 不能和　onnx包同时加载，即使同时加载了，也要onnx在torch.onnx之前进行加载。不然会报“段错误”，因此ｋｅ将torch->onnx以及onnx->tvm分开在两个py文件中
from pynvml import *     #pip install nvidia-ml-py3,显存监控
# from models_class.pose_resnet import * #关键点
from models_class.resnet_chejian_xiaolaomei import *

# from msra_resnet import * #centernet转换


# 将torch模型转为onnx模型

#转换关键点模型pth->onnx
def gen_torch_to_onnx_model_keypoints(model_path, save_onnx_path):
    # 加载torch模型，分类模型
    # 关键点模型
    model = get_pose_net('',False)  # pytorch的网络
    state_dict = torch.load(model_path)
    model.load_state_dict(state_dict)

    # 以下操作是避免pytorch转libtorch的模型显存增加
    for p in model.parameters():
        p.requires_grad = False
    model.cuda()
    model.eval()

    # 转 ONNX 模型
    example = torch.rand(32, 3, 1024, 768).cuda()   #输入图像大小,自行设定,如果使用tensorrt7.0.0多batch使用时，此处必须设定batch的大小

    input_name = ['input_1']  # 如果使用tensorrt7.0.0，必须设置输入的名字，输入可以是一个，也可以是多个,在c++中使用TRT需要知道的
    output_name = ['output_1']  # 如果使用tensorrt7.0.0，必须设置输出层的名字，输出可以是一个，也可以是多个,在c++中使用TRT需要知道的

    torch_out = torch.onnx.export(model=model,  # model being run
                                  args=(example,),  # model input (or a tuple for multiple inputs)
                                  f=save_onnx_path,
                                  input_names=input_name,
                                  output_names=output_name,
                                  export_params=True,
                                  verbose=True)  # store the trained parameter weights inside the model file   # 带参数输出)
    print("onnx convert ok")

#转换centernet检测模型pth->onnx
def gen_torch_to_onnx_model_centernet(model_path,save_onnx_path):

    #加载torch模型，分类模型

    #centernet
    num_layers = 10
    heads = {'hm': 7, 'reg': 2, 'wh': 2}
    head_conv = 64
    resnet_spec = {10: (BasicBlock, [1, 1, 1, 1]),
                   18: (BasicBlock, [2, 2, 2, 2]),
                   34: (BasicBlock, [3, 4, 6, 3]),
                   50: (Bottleneck, [3, 4, 6, 3]),
                   101: (Bottleneck, [3, 4, 23, 3]),
                   152: (Bottleneck, [3, 8, 36, 3])}
    name = 'ResNet_10'
    block_class, layers = resnet_spec[num_layers]
    model = PoseResNet(block_class, layers, heads, head_conv=head_conv)

    state_dict = torch.load(model_path)
    model.load_state_dict(state_dict['state_dict'])

    model.eval().cuda()

    # 转 ONNX 模型
    example = torch.rand(16, 3, 1120, 1120).cuda() #输入图像大小,自行设定,如果使用tensorrt7.0.0多batch使用时，此处必须设定batch的大小

    input_name = ['input_1'] #如果使用tensorrt7.0.0，必须设置输入的名字，输入可以是一个，也可以是多个
    output_name = ['output_hm','output_wh','output_reg'] #如果使用tensorrt7.0.0，必须设置输出层的名字，输出可以是一个，也可以是多个


    torch_out = torch.onnx.export(model=model,  # model being run
                                args=(example,),  # model input (or a tuple for multiple inputs)
                                f=save_onnx_path,
                                input_names=input_name,
                                output_names=output_name,
                                export_params=True,
                                verbose=True)  # store the trained parameter weights inside the model file   # 带参数输出)
    print("onnx convert ok")

#转换分类检测模型pth->onnx
def gen_torch_to_onnx_model_class(model_path,save_onnx_path):

    #加载torch模型，分类模型
    model = resnet29_cls() #pytorch的网络

    state_dict = torch.load(model_path)
    model.load_state_dict(state_dict,strict=True)

    model.eval().cuda()

    # 转 ONNX 模型
    example = torch.rand(32, 3, 400, 400).cuda() #输入图像大小,自行设定
    # example = torch.rand(2, 3, 608, 608).cuda() #输入图像大小,自行设定,如果使用tensorrt7.0.0多batch使用时，此处必须设定batch的大小

    input_name = ['input_1'] #如果使用tensorrt7.0.0，必须设置输入的名字，输入可以是一个，也可以是多个
    output_name = ['output_1'] #如果使用tensorrt7.0.0，必须设置输出层的名字，输出可以是一个，也可以是多个

    torch_out = torch.onnx.export(model=model,  # model being run
                                args=(example,),  # model input (or a tuple for multiple inputs)
                                f=save_onnx_path,
                                input_names=input_name,
                                output_names=output_name,
                                export_params=True,
                                verbose=True)  # store the trained parameter weights inside the model file   # 带参数输出)
    print("onnx convert ok")

def gen_torch_to_onnx_model(model_path,save_onnx_path):

    #加载torch模型，分类模型
    #关键点模型
    model = get_pose_net() #pytorch的网络

    #分类模型
    # model = resnet29_cls() #pytorch的网络

    #centernet
    # num_layers = 10
    # heads = {'hm': 7, 'reg': 2, 'wh': 2}
    # head_conv = 64
    # resnet_spec = {10: (BasicBlock, [1, 1, 1, 1]),
    #                18: (BasicBlock, [2, 2, 2, 2]),
    #                34: (BasicBlock, [3, 4, 6, 3]),
    #                50: (Bottleneck, [3, 4, 6, 3]),
    #                101: (Bottleneck, [3, 4, 23, 3]),
    #                152: (Bottleneck, [3, 8, 36, 3])}
    # name = 'ResNet_10'
    # block_class, layers = resnet_spec[num_layers]
    # model = PoseResNet(block_class, layers, heads, head_conv=head_conv)

    # state_dict = torch.load(model_path)
    # model.load_state_dict(state_dict['state_dict'])

    state_dict = torch.load(model_path)
    model.load_state_dict(state_dict,strict=True)

    model.eval().cuda()

    # 转 ONNX 模型
    example = torch.rand(32, 3, 200, 200).cuda() #输入图像大小,自行设定
    # example = torch.rand(2, 3, 608, 608).cuda() #输入图像大小,自行设定,如果使用tensorrt7.0.0多batch使用时，此处必须设定batch的大小

    # input_name = ['input_1'] #如果使用tensorrt7.0.0，必须设置输入的名字，输入可以是一个，也可以是多个
    # output_name = ['output_hm','output_wh','output_reg'] #如果使用tensorrt7.0.0，必须设置输出层的名字，输出可以是一个，也可以是多个


    input_name = ['input_1'] #如果使用tensorrt7.0.0，必须设置输入的名字，输入可以是一个，也可以是多个
    output_name = ['output_1'] #如果使用tensorrt7.0.0，必须设置输出层的名字，输出可以是一个，也可以是多个

    torch_out = torch.onnx.export(model=model,  # model being run
                                args=(example,),  # model input (or a tuple for multiple inputs)
                                f=save_onnx_path,
                                input_names=input_name,
                                output_names=output_name,
                                export_params=True,
                                verbose=True)  # store the trained parameter weights inside the model file   # 带参数输出)
    print("onnx convert ok")

if __name__=="__main__":
    # 显存监控
    GPU_USE=0
    nvmlInit() #初始化
    handle = nvmlDeviceGetHandleByIndex(GPU_USE) #获得指定GPU的handle
    info_begin = nvmlDeviceGetMemoryInfo(handle) #获得显存信息

    # 路径设置
    img_str = "/data_1/Working/project/torch-onnx-tvm/torch_py/data/000d81205274490bb1729e4e26226f01-3-1@00.jpg"
    model_path = "/data1/Working/project/weifa-all-project/video-3.0/use-tensorRT/test_centernet/models/xiang/centerNet.pth"
    onnx_save_path = "/data1/Working/project/weifa-all-project/video-3.0/use-tensorRT/test_centernet/models/xiang/centerNet_16_1120_1120.onnx"

    # 将torch模型转为onnx模型
    gen_torch_to_onnx_model_centernet(model_path,onnx_save_path)
    # gen_torch_to_onnx_model_keypoints(model_path,onnx_save_path)
    # gen_torch_to_onnx_model_class(model_path,onnx_save_path)

    # #输出显存使用mb
    info_end = nvmlDeviceGetMemoryInfo(handle)
    print("-" * 15 + "TORCH GPU MEMORY INFO" + "-" * 15)
    print("       Memory Total: " + str(info_end.total // (1024 ** 2)))
    print("       Memory Free: " + str(info_end.free // (1024 ** 2)))
    print("       Memory Used: " + str(info_end.used // (1024 ** 2) - info_begin.used // (1024 ** 2)))
    print("-" * 40)
IT捕快
关注
  • 0
    点赞
  • 踩
  • 5
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
专栏目录
pytorch训练图像分类模型pth转ONNX并测试
10-29
1. 搭建自己的简单二分类网络，使用pytorch训练和测试； 2. 将pytorch训练的pth模型转换成ONNX，并编码测试； 3. 含训练和测试数据，含训练ok的pth模型和ONNX模型，含完整python和C++实现； 4. 使用方法：首先运行“TrainTestConvertOnnx.py”执行“训练数据读入、模型训练、模型测试、导出onnx”，再运行“TestOnnx.cpp”测试onnx（需要配置OpenCV）；
torch.onnx.export():将pytorch模型转换为.onnx模型
热门推荐
QFJIZHI的博客
04-01 4万+
torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None, aten=False, export_raw_ir=False, operator_export_type=None, opset_version=Non...
模型训练----将pth模型转换为onnx
最新发布
混吃等死，研究生在读，毕业无望，偶尔搞点好玩的
03-28 879
需要在虚拟环境中安装onnx和onnxruntime（GPU），环境和自己的cuda版本要对应上激活环境，查看环境的cuda版本,我是cuda11.6 +cudnn8302，那就选择1.14吧输入以下指令进行安装环境。
torch转onnx
cjk_cynosure的博客
11-10 201
torch转onnx。
torch转onnx模型
Y_snower的博客
06-24 3861
以分类为例torch转换为onnx模型
torch-1.0.0.rar
06-10
torch-1.0.0.rar
onnx-simplifier:简化您的Onnx模型
04-29
ONNX简化器 ONNX很棒，但有时太复杂了。 背景 有一天，我想将以下简单的重塑操作导出到ONNX： import torch class JustReshape ( torch ....model_name = 'just_reshape.onnx' dummy_input = torch . randn ( 2 ,
yolov5-onnx-tensorrt:如何使用 TensorRT 运行 yolov5 模型
05-29
请使用torch>=1.6.0 + onnx==1.8.0 + TensorRT 7.0.0.11 运行代码 代码结构 networks代码是网络demo代码在 Jetson AGX Xavier 上运行 tensorrt 实现 ├── networks ├── utils ├── models ├── demo │ ├...
go-torch-master.rar
05-24
go-torch全量包 github.com/fatih github.com/jessevdk github.com/stretchr
torch-1.4.0-百度云下载地址.txt
03-28
ubuntu16.04 64位系统 python2.7 torch安装包torch-1.4.0-cp27-cp27mu-manylinux1_x86_64.whl
torch 网络模型转换onnx格式，并可视化
桜
06-08 1188
1. 网络模型结构 2. 转为onnx格式 3. 通过netron查看网络结构 3.1 netron安装 3.2 netron可视化
将torch转换为onnx的脚本
Cheny的博客
08-24 424
将torch模型转化为onnx模型的脚本
Pytorch模型(.pth)转onnx模型(.onnx)
浩瀚之水的专栏
09-11 1755
简介 Open Neural Network Exchange（ONNX，开放神经网络交换）格式，是一个用于表示深度学习模型的标准，可使模型在不同框架之间进行转移。 ONNX是一种针对机器学习所设计的开放式的文件格式，用于存储训练好的模型。它使得不同的人工智能框架（如Pytorch, MXNet）可以采用相同格式存储模型数据并交互。 ONNX的规范及代码主要由微软，亚马逊 ，Facebook 和 IBM等公司共同开发，以开放源代码的方式托管在Github上。目前官方支持加载ONNX模型并进行推理的深度学习
pytorch pth模型转onnx
qq_43994782的博客
11-25 1495
转换方法： 使用pytorch自带转换函数转换 import torch # 实例你的模型 net = MyModel() print('load static dict') # 加载参数！ （如果忘记的话，，，就重来呗） net.load_state_dict(torch.load("net.pth", map_location=torch.device('cpu'))) net.eval().cuda() print('to cuda') input1 = torch.randn(1, 3, 112
使用onnx包将pth文件转换为onnx文件
qq_45893319的博客
12-24 3318
本文对比一下两种pth文件转为onnx的区别以及onnx文件在NETRON中的图
Pytorch 转ONNX详解
陶将的博客
11-06 6438
模型部署面临两大问题，如何让开发者不受开发框架的限制和如何提升模型的运算效率。经过工业界和学术界数年的探索，一般流行的模型部署流水线是：首先开发者们在任意一种深度学习框架下定义网络结构，训练得到网络参数；然后，将模型的结构和参数转换成一种只描述网络结构的中间表示。针对网络结构的优化会在这步进行。ONNX就是一种模型中间表示。 最后，用面向硬件的高性能编程框架(如 CUDA，OpenCL）编写能高效执行深度学习网络中算子的推理引擎，推理引擎将中间表示转换成特定的文件格式。
将训练好的pytorch模型的pth文件转换成onnx模型（亲测成功）
hutao1030813002的博客
03-18 1万+
将训练好的pytorch模型的pth文件转换成onnx模型（亲测成功） 模型转换 由于我们要进行openvino部署，pytorch模型不能直接转化为openvino部署所需要的IR中间文件。 openvino只接受tensorflow、caffe、onnx等模型直接生成IR文件。因此我们要先将pytorch模型转化为onnx模型，在将onnx模型生成.xml和.bin文件。 代码如下： import torch import torchvision dummy_input = torch.randn(1
Pytorch转ONNX详解
zandaoguang的博客
12-10 1万+
作者丨立交桥跳水冠军@知乎来源丨https://zhuanlan.zhihu.com/p/272767300编辑丨极市平台之前几个月参与了OpenMMlab的模型转ONNX的工作（git...
torch.onnx.export导出onnx,支持多个模型输入
05-23
是的，torch.onnx.export函数可以将PyTorch模型导出为ONNX格式的模型，支持多个模型输入。在导出时，可以通过传递一个元组作为输入参数来指定多个输入。例如： ```python import torch class MyModel(torch.nn.Module): def __init__(self): super().__init__() self.conv1 = torch.nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1) self.conv2 = torch.nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1) self.fc1 = torch.nn.Linear(128 * 8 * 8, 512) self.fc2 = torch.nn.Linear(512, 10) def forward(self, x): x = torch.nn.functional.relu(self.conv1(x)) x = torch.nn.functional.max_pool2d(x, 2) x = torch.nn.functional.relu(self.conv2(x)) x = torch.nn.functional.max_pool2d(x, 2) x = x.view(-1, 128 * 8 * 8) x = torch.nn.functional.relu(self.fc1(x)) x = self.fc2(x) return x model = MyModel() # 保存模型为ONNX格式 dummy_input_1 = torch.randn(1, 3, 32, 32) dummy_input_2 = torch.randn(1, 3, 32, 32) output_file_path = "my_model.onnx" torch.onnx.export(model, (dummy_input_1, dummy_input_2), output_file_path) ``` 在这个例子中，我们定义了一个包含两个输入的简单的卷积神经网络，并将其导出为ONNX格式。我们使用`torch.randn`创建两个虚拟输入，并将它们作为元组传递给导出函数。在导出的ONNX文件中，我们将看到两个输入节点，一个名为`input.1`，另一个名为`input.2`。

“相关推荐”对你有帮助么？

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值