torch转onnx模型加速,onnx模型推理直接达到tensorRT速度,省去onnx转tensorRT的流程

完成torch转onnx后,直接基于onnx模型推理达到tensorRT速度~


最近在进行torch模型推理加速过程时,原本想的方案是torch转onnx再转tensorRT实现加速,但是在转完onnx,使用onnx模型验证推理效果时发现,可以直接通过设置onnx runtime 推理模型的provider参数,实现tensorRT的推理速度。

在未设置provider时速度很慢,我这默认应该用的是CPU。

通过以下方式修改为tensorRT引擎:

import onnxruntime as ort

onnx_model = ort.InferenceSession(landmark_gen_onnx_path, providers=['TensorrtExecutionProvider'])

我这边是直接就运行成功了,并且后面通过转完tensorRT模型后的速度验证,发现这种方式下比我自己转tensorRT模型再用tensorRT模型推理速度还快一点点。

provider 参数解释:
providers: Optional sequence of providers in order of decreasing precedence. Values can either be provider names or tuples of (provider name, options dict). If not provided, then all available providers are used with the default precedence.

可选参数为:‘TensorrtExecutionProvider’, ‘CUDAExecutionProvider’, ‘CPUExecutionProvider’,大概分别代表使用TensorRT,CUDA,CPU 执行。

  • 13
    点赞
  • 4
    收藏
    觉得还不错? 一键收藏
  • 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. 使用TensorRTONNX模型优化为序列化引擎: ``` 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、付费专栏及课程。

余额充值