神经网络模型模型转ONNX

近期由于业务需要,需要将训练好的模型转为ONNX格式,为此颇费了一番功夫,在此总结一下吧。。

1、ONNX是一种神经网络模型保存的中间格式,支持多种格式的模型转为ONNX,也支持使用ONNX导入多种格式的模型,具体见https://github.com/onnx/tutorials;目前其实ONNX对于模型的支持还不是太好,主要表现在一些op还不能够支持;

2、在PyTorch下要将模型保存成ONNX格式需要使用torch.onnx.export()函数,使用该函数的时候需要传入下面参数:

  --model:待保存的model,也就是你在程序中已经训练好或者初始化好的模型

  --input_shape:指定输入数据的大小,也就是输入数据的形状,是一个包含输入形状元组的列表;

  --name:模型的名称,即模型的保存路径;

  --verbrose:True或者False,用来指定输出模型时是否将模型的结构打印出来;

  --input_names:输入数据节点的名称,数据类型为包含字符串的列表;一般将这个名称设为['data'];

  --output_names:输出数据节点的名称,类型与输入数据的节点名称相同;

  在成功导出模型后,可以使用ONNX再对模型进行检查:

import onnx

# Load the ONNX model
model = onnx.load("alexnet.onnx")

# Check that the IR is well formed
onnx.checker.check_model(model)

# Print a human readable representation of the graph
onnx.helper.printable_graph(model.graph)

目前PyTorch还不支持导入ONNX格式的模型。

3、使用MXNET导出模型为ONNX时,参考地址:https://cwiki.apache.org/confluence/display/MXNET/ONNXhttp://mxnet.incubator.apache.org/versions/master/tutorials/onnx/export_mxnet_to_onnx.html。MXNet模型的保存格式为.json文件+.params文件,.json文件里保存的是模型的结构,.params文件中保存的是模型的参数。使用onnx_mxnet.export_export_model()方法就可以实现将模型从mxnet转为ONNX格式,该方法需要传入的参数为:

  --sym:.json文件,也就是保存了网络结构的文件

  --params:参数文件

  --input_shape:输入数据的形状,是一个包含形状元组的列表

  --input_type:输入数据的类型;

  --模型的保存路径

4、从MXNet导入ONNX格式模型:需要使用mxnet.contrib.onnx.onnx2mx.import_model.import_model(model_file),这里返回的是sym, arg_arams,aux_params,也就是网络结构symbol对象,保存参数的字典, 再将其转为MXNet的module对象(使用mxnet.module.Module()),即可将模型恢复到mxnet框架下可执行的模型。

 

最后,好久没有记录日常学习积累的东西了,趁着失眠开个好头吧,晚安。。。

转载于:https://www.cnblogs.com/puheng/p/10873289.html

  • 2
    点赞
  • 5
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
是的,TensorRT可以将PyTorch中的神经网络模型换为ONNX格式。TensorRT提供了一个Python API,您可以使用它来将PyTorch模型换为ONNX格式,然后使用TensorRT将其优化为适用于GPU加速推理的序列化引擎。具体步骤如下: 1. 将PyTorch模型换为ONNX格式: ``` import torch import onnx # Load the PyTorch model model = torch.load('model.pt') # Convert the PyTorch model to ONNX dummy_input = torch.randn(1, 3, 224, 224) input_names = ['input'] output_names = ['output'] onnx_path = 'model.onnx' torch.onnx.export(model, dummy_input, onnx_path, verbose=False, input_names=input_names, output_names=output_names) ``` 2. 使用TensorRT将ONNX模型优化为序列化引擎: ``` import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit # Load the ONNX model onnx_path = 'model.onnx' onnx_model = onnx.load(onnx_path) # Create a TensorRT builder and network trt_logger = trt.Logger(trt.Logger.WARNING) trt_builder = trt.Builder(trt_logger) trt_network = trt_builder.create_network() # Create an ONNX parser to parse the ONNX model into the TensorRT network onnx_parser = trt.OnnxParser(trt_network, trt_logger) onnx_parser.parse(onnx_model.SerializeToString()) # Set the maximum batch size and maximum workspace size trt_builder.max_batch_size = 1 trt_builder.max_workspace_size = 1 << 30 # Build the TensorRT engine from the TensorRT network trt_engine = trt_builder.build_cuda_engine(trt_network) # Serialize the TensorRT engine to a file trt_engine_path = 'model.engine' with open(trt_engine_path, 'wb') as f: f.write(trt_engine.serialize()) ``` 3. 使用TensorRT引擎进行推理: ``` import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit import numpy as np # Load the serialized TensorRT engine trt_engine_path = 'model.engine' with open(trt_engine_path, 'rb') as f: trt_engine_data = f.read() # Create a TensorRT runtime and deserialize the TensorRT engine trt_logger = trt.Logger(trt.Logger.WARNING) trt_runtime = trt.Runtime(trt_logger) trt_engine = trt_runtime.deserialize_cuda_engine(trt_engine_data) # Create a TensorRT execution context trt_context = trt_engine.create_execution_context() # Allocate GPU memory for the input and output tensors input_shape = (1, 3, 224, 224) output_shape = (1, 1000) input_dtype = np.float32 output_dtype = np.float32 input_size = np.product(input_shape) * np.dtype(input_dtype).itemsize output_size = np.product(output_shape) * np.dtype(output_dtype).itemsize input_gpu = cuda.mem_alloc(input_size) output_gpu = cuda.mem_alloc(output_size) # Create a CUDA stream stream = cuda.Stream() # Initialize the input tensor with random data input_cpu = np.random.rand(*input_shape).astype(input_dtype) cuda.memcpy_htod_async(input_gpu, input_cpu, stream) # Run inference on the TensorRT engine trt_context.execute_async(1, [int(input_gpu), int(output_gpu)], stream.handle, None) # Copy the output tensor back to the CPU output_cpu = np.empty(output_shape, dtype=output_dtype) cuda.memcpy_dtoh_async(output_cpu, output_gpu, stream) # Synchronize the CUDA stream stream.synchronize() # Print the output tensor print(output_cpu) ```

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

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

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值