TensorRT 原理及核心代码示例

        TensorRT 是 NVIDIA 推出的高性能 深度学习推理优化器，通过模型压缩、内核自动调优、动态内存优化等技术，显著提升模型在 GPU 上的推理速度。以下是其核心原理及代码实现。

一、TensorRT 核心原理

1. 优化流程

• 步骤 1：解析模型
  将训练好的模型（如 PyTorch/TensorFlow）转换为 TensorRT 支持的格式（如 ONNX）。

• 步骤 2：优化阶段

  • 层融合（Layer Fusion）：合并多个计算层（如 Conv + BN + ReLU）为单一内核，减少内存访问。

  • 精度校准（Precision Calibration）：支持 FP32/FP16/INT8 精度，自动选择最优计算模式。

  • 内核自动调优（Kernel Auto-Tuning）：根据 GPU 架构生成高效 CUDA 内核。

• 步骤 3：生成引擎（Engine）
  将优化后的模型序列化为可执行的推理引擎（.plan 或 .trt 文件）。

2. 核心组件

• Builder：构建优化引擎的核心接口。

• Network：定义模型计算图的抽象层。

• Parser：解析原始模型（如 ONNX）到 TensorRT 网络。

• ExecutionContext：执行推理的上下文环境。

二、TensorRT 核心代码示例

示例场景：将 PyTorch 模型转换为 TensorRT 引擎并推理

步骤 1：导出 PyTorch 模型为 ONNX

import torch
import torchvision.models as models

# 加载预训练 ResNet-50
model = models.resnet50(pretrained=True)
model.eval()

# 输入示例（动态 batch 和分辨率）
dummy_input = torch.randn(1, 3, 224, 224)  # (batch, channels, H, W)

# 导出 ONNX 模型
torch.onnx.export(
    model,
    dummy_input,
    "resnet50.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size", 2: "height", 3: "width"}},
)

步骤 2：使用 TensorRT 构建优化引擎

import tensorrt as trt

# 初始化 TensorRT 组件
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

# 解析 ONNX 模型
with open("resnet50.onnx", "rb") as f:
    parser.parse(f.read())

# 配置优化参数
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16)  # 启用 FP16 加速
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB 显存

# 构建引擎
engine = builder.build_engine(network, config)

# 序列化引擎到文件
with open("resnet50.engine", "wb") as f:
    f.write(engine.serialize())

步骤 3：执行推理

import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

# 反序列化引擎
with open("resnet50.engine", "rb") as f:
    runtime = trt.Runtime(logger)
    engine = runtime.deserialize_cuda_engine(f.read())

# 创建执行上下文
context = engine.create_execution_context()

# 分配输入输出内存
inputs, outputs, bindings = [], [], []
stream = cuda.Stream()

for binding in engine:
    size = trt.volume(engine.get_binding_shape(binding))  # 计算内存大小
    dtype = trt.nptype(engine.get_binding_dtype(binding))
    host_mem = cuda.pagelocked_empty(size, dtype)
    device_mem = cuda.mem_alloc(host_mem.nbytes)
    bindings.append(int(device_mem))
    if engine.binding_is_input(binding):
        inputs.append({"host": host_mem, "device": device_mem})
    else:
        outputs.append({"host": host_mem, "device": device_mem})

# 准备输入数据（假设输入为 224x224 图像）
np_input = np.random.randn(1, 3, 224, 224).astype(np.float32)
np.copyto(inputs[0]["host"], np_input.ravel())

# 执行推理
cuda.memcpy_htod_async(inputs[0]["device"], inputs[0]["host"], stream)
context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
cuda.memcpy_dtoh_async(outputs[0]["host"], outputs[0]["device"], stream)
stream.synchronize()

# 获取输出结果
output = outputs[0]["host"].reshape(1, 1000)  # ResNet-50 输出为 1000 类
print("推理结果形状:", output.shape)

三、TensorRT 优化效果

优化技术	加速比	适用场景
FP32 → FP16	2-3x	大多数 GPU 支持
FP32 → INT8	3-5x	需要校准（需代表性数据集）
层融合	1.5-2x	复杂模型（如 ResNet）

四、关键注意事项

1.动态形状支持
使用 dynamic_axes 定义可变维度（如 batch 或分辨率）：

        profile = builder.create_optimization_profile()
        profile.set_shape("input", (1,3,224,224), (4,3,224,224), (8,3,224,224))
        config.add_optimization_profile(profile)

2.INT8 校准
需提供校准数据集生成量化参数：

config.set_flag(trt.BuilderFlag.INT8)
calibrator = trt.Int8EntropyCalibrator2(calibration_data)
config.int8_calibrator = calibrator

3.插件支持
自定义算子需通过 TensorRT 插件实现（如检测模型中的非标准层）。

TensorRT 通过 内核优化、内存复用、混合精度计算 将 GPU 推理性能推向极限。其核心思想是 静态优化 + 运行时高效执行，适用于对延迟敏感的部署场景（如自动驾驶、实时视频分析）。掌握 TensorRT 的优化技巧，可将模型推理速度提升 2-5 倍。