TensorRT python版本 windows上下载安装：
TensorRT python Windows 下载安装

写在前面：

经过阅读多篇文章、查看官方文档、论坛，终于成功打通了python版本的TensorRT流程。
本文介绍的是python版本的tensorRT推理流程，测试pytroch框架保存的分类模型。本人综合测试推理速度比pytorch快一倍以上，该数据仅供参考。

python==3.9.20 onnx==1.17.0 numpy==1.24.4 tensorrt==10.6.0

整体流程：

头文件

import tensorrt as trt
import torch
import pycuda.driver as cuda
import pycuda.autoinit
from mobilenetv3 import MobileNetV3_Large  #分类网络
import numpy as np
import cv2

1.根据pytorch模型生成onnx模型
需要安装onnx库。

#推荐使用加载模型参数的方法。
model = torch.load('model.pth')
# 切换到评估模式，禁用 Dropout 等层
model.eval()  
# 根据模型的输入形状调整
dummy_input = torch.randn(1, 3, 112, 112).to(device)
onnx_model_path = "model.onnx"
torch.onnx.export(model, dummy_input, onnx_model_path,
                  input_names=['input'], output_names=['output'],
                  opset_version=11,
                  dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})

2.根据onnx模型生成engine

# 创建logger：日志记录器
logger = trt.Logger(trt.Logger.WARNING)
# 创建构建器builder
builder = trt.Builder(logger)
# 预创建网络
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
# 加载onnx解析器
parser = trt.OnnxParser(network, logger)
success = parser.parse_from_file('model.onnx')

# builder配置
config = builder.create_builder_config()

profile = builder.create_optimization_profile()
input_tensor = network.get_input(0)
# 定义输入的动态维度，最小、最大和理想形状
min_shape = (1, 3, 112, 112)  # 假设最小的形状
opt_shape = (16, 3, 112, 112)  # 假设理想的形状
max_shape = (64, 3, 112, 112)  # 假设最大的形状
profile.set_shape(input_tensor.name, min_shape, opt_shape, max_shape)
config.add_optimization_profile(profile)
# 分配显存作为工作区间，一般建议为显存一半的大小
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)
serialized_engine = builder.build_serialized_network(network, config)
# 序列化生成engine文件
with open('model.engine', "wb") as f:
    f.write(serialized_engine)
    print("generate file success!")

3.加载engine进行推理

# 使用 trt.Runtime 加载 engine
with open("model.engine", "rb") as f:
    engine_data = f.read()

# 创建 Runtime 对象
runtime = trt.Runtime(logger)
# 反序列化引擎
engine = runtime.deserialize_cuda_engine(engine_data)
#输入输出形状，图像大小为112*112 9分类任务
input_shape=(1,3,112,112)
output_shape=(1,9)
#读取图像
input_image_path= 'image.png'
image=cv2.imread(input_image_path)
image=cv2.cvtColor(image,cv2.COLOR_BGR2RGB).astype(np.float32)/255.0
image=cv2.resize(image,(112,112))
# 从 (112, 112, 3) 转换为 (3, 112, 112)
image = np.transpose(image, (2, 0, 1))  
image=image.copy()
# 增加一个批量维度，得到 (1, 3, 112, 112)
image_data = np.expand_dims(image, axis=0)
with engine.create_execution_context() as context:
    # 分配CPU锁页内存和GPU显存
    h_input = cuda.pagelocked_empty(input_shape, dtype=np.float32)
    h_output = cuda.pagelocked_empty(output_shape, dtype=np.float32)
    # 分配设备内存
    d_input = cuda.mem_alloc(h_input.nbytes)
    d_output = cuda.mem_alloc(h_output.nbytes)
    # 输入输出缓冲区
    context.set_tensor_address(engine.get_tensor_name(0), int(d_input))  
    context.set_tensor_address(engine.get_tensor_name(1), int(d_output))  
    # 创建cuda流
    stream = cuda.Stream()
    # 将图像数据复制到页面锁定内存
    np.copyto(h_input, image_data)
    # 将数据拷贝到设备内存
    cuda.memcpy_htod_async(d_input, h_input,stream)
    #执行推理
    context.execute_async_v3(stream_handle = stream.handle)
    # 将输出数据从设备内存拷贝回主机内存
    cuda.memcpy_dtoh_async(h_output, d_output, stream)
    print(h_output)

总结

用户在使用 TensorRT 时可以通过以下设置来调节推理性能：
精度优化（FP16、INT8）
最大批处理大小
工作空间大小
优化器策略（优化配置）
CUDA 流管理
层融合和优化
自定义插件
执行上下文管理
根据具体的硬件环境和推理需求，用户可以调节这些参数来实现最佳的性能。调整时需要在性能和准确性之间找到合适的平衡。