TensorRT(7)：python版本使用入门

最新推荐文章于 2025-03-26 10:17:22 发布

hjxu2016

最新推荐文章于 2025-03-26 10:17:22 发布

阅读量2.8k

点赞数 3

分类专栏：推理框架|TensorRT

本文链接：https://blog.youkuaiyun.com/hjxu2016/article/details/115721812

版权

推理框架|TensorRT 专栏收录该内容

12 篇文章

订阅专栏

本文详细介绍了如何使用TensorRT的Python API在线加载ONNX模型，序列化保存为TRT模型，以及反序列化加载已保存的模型进行推断。通过示例代码展示了从读取图像、模型转换到执行推断的完整流程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

TensorRT系列传送门(不定期更新): 深度框架|TensorRT

文章目录

本文以onnx分类模型为例，简单介绍TRT-python版本的使用流程，这里不涉及量化，就以fp32为例，验证在线加载onnx模型、序列化保存为trt模型、加载序列后的模型等。

TensorRT的使用过程包括两个阶段， build和 deployment
在这里插入图片描述
build阶段主要完成模型转换，从caffe或者tf或者pyTorch到TensorRT

一、TRT在线加载模型，并构建推断引擎的步骤

主要分为四步：

1、建立一个logger日志，必须要有，但又不是那么重要
2、创建一个builder
3、创建一个netwok，这时候netWork只是一个空架子
4、建立一个 Parser，caffe模型、onnx模型和TF模型都有对应的paser，顾名思义，就用用来解析模型文件的
5、建立 engine，进行层之间融合或者进度校准方式，可以fp32、fp16或者fp8
6、建议一个context，这个是用来做inference推断的。上面连接engine，下对应推断数据，所以称之为上下文联系器(博主自取的)
(PS:感谢晖神提供流程图)

二、在线加载onnx模型,序列化保存到本地

主要功能：
1、opencv读取图像，减均值除方差，在NHWC转NCHW
2、trt解析onnx模型
3、保存序列化后的模型
4、建立engine，推断模型、生成结果

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""                  
*  * *** *  * *  *      
*  *  *   **  *  *             
****  *   **  *  *                 
*  *  *   **  *  *         
*  * **  *  * ****  

@File     :inferenceWithTrt.py  
@Date     :2021/3/18 下午1:18  
@Require  :   
@from     :https://blog.youkuaiyun.com/hjxu2016/
@Funtion  :用trt接口推断onnx模型
"""

import pycuda
import pycuda.driver as cuda
import pycuda.autoinit
import onnx
import onnxruntime
import tensorrt as trt
import cv2
import numpy as np
 
# 为输入输出分配内存

def read_img():
    bgr_img = cv2.imread('./Data/validation/0-back/dog.4.jpg')

    bgr_img = cv2.resize(bgr_img, (224, 224))
    rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)
    # img = Image.open("./Data/validation/0-back/dog.4.jpg")
    # img = img.resize((224, 224))
    x = np.transpose(rgb_img, (2, 0, 1))
    x = x/255.
    return x

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
onnx_model_file = "./saved_weights/senet_50.onnx"
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
builder = trt.Builder(TRT_LOGGER)  # 创建一个builder
config = builder.create_builder_config() # 创建一个congig
config.max_workspace_size = 1<<20
network = builder.create_network(EXPLICIT_BATCH) # 创建一个network
parser = trt.OnnxParser(network, TRT_LOGGER)

model = open(onnx_model_file, 'rb')
if not parser.parse(model.read()):
    for error in range(parser.num_errors):
        print(parser.get_error(error))

engine = builder.build_cuda_engine(network) # 创建engine

# 是否需要序列化保存模型
if False:
    saved_trt_path = "./saved_weights/senet_50.trt"
    with open(saved_trt_path, "wb") as f:
        f.write(engine.serialize())
# 构建context
context = engine.create_execution_context() # 创建context用来执行推断

三、反序列化直接加载保存后的trt模型

saved_trt_path = "./saved_weights/senet_50.trt"
f = open(saved_trt_path, "rb") 
runtime = trt.Runtime(TRT_LOGGER) 
engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context() # 创建context用来执行推断

四、构建cuda数据，推断trt模型

# h指设备端数据，d指的是cuda端数据
h_input = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)), dtype=trt.nptype(trt.float32))
h_output = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)), dtype=trt.nptype(trt.float32))
d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)
stream = cuda.Stream()
#读取数据
x = read_img().astype(trt.nptype(trt.float32)).ravel()
np.copyto(h_input, x)
cuda.memcpy_htod_async(d_input, h_input, stream)
context.execute_async(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
# 从GPU中返回output。
cuda.memcpy_dtoh_async(h_output, d_output, stream)
# 同步流。
stream.synchronize()
print(h_output)