深入TensorRT推理：Python实现与高效代码实践

一、TensorRT推理概述

TensorRT是NVIDIA推出的高性能深度学习推理引擎，专为优化和加速生产环境中的深度学习模型而设计。其核心优势在于通过层融合、精度校准、内核自动选择等技术，显著提升模型在NVIDIA GPU上的推理速度，同时降低延迟和内存占用。在Python生态中，TensorRT通过tensorrt Python包提供编程接口，结合ONNX等模型交换格式，实现了从训练到部署的无缝衔接。

1.1 TensorRT的核心优势

• 性能优化：通过图优化、层融合等技术，减少计算冗余，提升吞吐量。
• 动态形状支持：支持输入尺寸的动态变化，适应多样化场景需求。
• 多精度推理：支持FP32、FP16、INT8等多种精度，平衡精度与速度。
• 跨平台兼容：与NVIDIA GPU架构深度集成，支持从嵌入式设备到数据中心的多场景部署。

1.2 Python在TensorRT推理中的角色

Python凭借其简洁的语法和丰富的生态，成为TensorRT推理开发的首选语言。通过tensorrt Python包，开发者可以：

• 加载预训练模型（如ONNX格式）。
• 构建优化引擎。
• 执行高效推理。
• 集成到现有Python应用中。

二、Python TensorRT推理环境配置

2.1 安装TensorRT

TensorRT的安装需根据操作系统和CUDA版本选择对应版本。以Ubuntu 20.04和CUDA 11.x为例：

# 添加NVIDIA仓库（需先安装CUDA和cuDNN）
sudo apt-get update
sudo apt-get install -y --no-install-recommends \
    libnvinfer8 \
    libnvonnxparsers8 \
    python3-libnvinfer-dev
# 通过pip安装tensorrt包（可选，推荐使用NVIDIA官方仓库）
pip install nvidia-tensorrt

验证安装：

import tensorrt as trt
print(f"TensorRT版本: {trt.__version__}")

2.2 依赖库准备

• ONNX：用于模型转换（pip install onnx）。
• PyCUDA：GPU计算加速（pip install pycuda）。
• NumPy：数据处理基础库。

三、TensorRT推理代码实现

3.1 模型转换：从ONNX到TensorRT引擎

将训练好的模型（如PyTorch）导出为ONNX格式，再转换为TensorRT引擎：

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import onnx
# 1. 创建TensorRT Logger
logger = trt.Logger(trt.Logger.INFO)
# 2. 创建Builder和Network
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
# 3. 创建ONNX Parser
parser = trt.OnnxParser(network, logger)
onnx_model_path = "model.onnx"
with open(onnx_model_path, "rb") as f:
    if not parser.parse(f.read()):
        for error in range(parser.num_errors):
            print(parser.get_error(error))
        raise RuntimeError("ONNX解析失败")
# 4. 配置Builder
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB工作空间
# 5. 构建引擎
engine = builder.build_engine(network, config)
# 6. 序列化引擎
with open("model.engine", "wb") as f:
    f.write(engine.serialize())

关键点：

• 显式批处理：使用EXPLICIT_BATCH标志明确批处理维度。
• 内存管理：通过set_memory_pool_limit控制内存使用。
• 错误处理：检查ONNX解析错误，确保模型兼容性。

3.2 推理执行：Python代码示例

加载引擎并执行推理：

def load_engine(engine_path):
    with open(engine_path, "rb") as f, trt.Runtime(logger) as runtime:
        return runtime.deserialize_cuda_engine(f.read())
def infer(engine, input_data):
    context = engine.create_execution_context()
    # 分配输入/输出缓冲区
    for binding in engine:
        if engine.binding_is_input(binding):
            input_shape = engine.get_binding_shape(binding)
            input_size = trt.volume(input_shape) * engine.max_batch_size
            d_input = cuda.mem_alloc(input_size * 4)  # FP32
        else:
            output_shape = engine.get_binding_shape(binding)
            output_size = trt.volume(output_shape) * engine.max_batch_size
            d_output = cuda.mem_alloc(output_size * 4)
    # 传输输入数据到GPU
    cuda.memcpy_htod_async(d_input, input_data.astype(np.float32), stream)
    # 执行推理
    context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
    # 传输输出数据回CPU
    cuda.memcpy_dtoh_async(output, d_output, stream)
    stream.synchronize()
    return output
# 使用示例
engine = load_engine("model.engine")
input_data = np.random.rand(1, 3, 224, 224).astype(np.float32)  # 示例输入
output = infer(engine, input_data)
print("推理结果:", output)

优化建议：

• 异步执行：使用execute_async_v2提升吞吐量。
• 流管理：通过pycuda.driver.Stream实现并行计算。
• 动态形状：若模型支持动态输入，需在构建引擎时指定范围。

四、性能优化与调试

4.1 精度校准与INT8量化

INT8量化可显著提升推理速度，但需校准以减少精度损失：

# 创建INT8校准器
class MyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, input_shapes, cache_file="calibration.cache"):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.cache_file = cache_file
        # 初始化输入数据生成逻辑
    def get_batch_size(self):
        return 1
    def get_batch(self, names):
        # 返回校准数据（需实现）
        pass
    def read_calibration_cache(self, size):
        # 读取缓存
        pass
    def write_calibration_cache(self, ptr, size):
        # 写入缓存
        pass
# 在Builder配置中启用INT8
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = MyCalibrator(input_shapes)

4.2 调试与日志

• 日志级别：通过trt.Logger设置日志级别（WARNING、INFO、VERBOSE）。
• 性能分析：使用trt.Profiler监控各层耗时。

五、实际应用建议

1. 模型兼容性：确保ONNX模型符合TensorRT支持的操作集。
2. 批量处理：合理设置批处理大小以最大化GPU利用率。
3. 持续集成：将TensorRT引擎生成纳入CI/CD流程，确保部署一致性。
4. 多引擎管理：对于多模型场景，使用trt.ICudaEngine池化资源。

六、总结

TensorRT通过Python接口为开发者提供了高效、灵活的深度学习推理解决方案。从模型转换到引擎构建，再到推理执行，每一步均需关注兼容性、性能和资源管理。通过结合INT8量化、异步执行等优化技术，可进一步提升推理效率。未来，随着TensorRT对更多模型和硬件的支持，其在边缘计算、自动驾驶等领域的应用将更加广泛。