TensorRT介绍与示例代码学习

下面的内容会分成几部分来讲：

1. 什么是 TensorRT（从宏观到细节）
2. 为什么需要 TensorRT（性能、部署和工程层面的原因）
3. 典型工作流程（从训练模型到上线推理）
4. 入门学习路线建议
5. 代码示例：ONNX → TensorRT → 推理 & 性能对比（Python）
6. 进一步的进阶方向与验证方法

---

1. 什么是 TensorRT？

用一句话概括：
TensorRT 是 NVIDIA 提供的、面向深度学习推理阶段的高性能优化与执行引擎 ，专门用来在 NVIDIA GPU 上做 高吞吐、低延迟 的推理。

更细一点地说，TensorRT 主要做几件事：

• 模型图优化 （Graph Optimization）
  • 合并算子（layer/operation fusion），减少内存访问和中间结果
  • 常量折叠、算子重排序等
• 精度优化 （Precision Calibration & Quantization）
  • 支持 FP32 / FP16 / INT8 等多种精度
  • 通过量化（尤其 INT8）大幅提升吞吐、减少显存占用
• 内存与执行规划 （Memory & Execution Planning）
  • 自动规划中间张量的复用，减少显存峰值
  • 为你的网络选择最合适的 kernel 实现和并行策略
• 设备相关优化
  • 针对不同 GPU 架构（Turing/Ampere/Ada/Hopper）选择最优 kernel
  • 支持 Tensor Cores 的加速（FP16/INT8 等）

从软件形态上看，TensorRT 主要包括：

• 一个 C++/Python API 库（核心是 ICudaEngine 和 IExecutionContext）
• 命令行工具（trtexec）
• 与框架的集成插件（Torch-TensorRT、TF-TRT、ONNX-TRT 等）
• 一堆高性能的算子实现（convolution, matmul, activation, etc.）

---

2. 为什么需要 TensorRT？

你已经有一个 "能跑起来" 的深度学习模型，为什么还要 TensorRT？

2.1 性能：更高的 QPS / 更低的延迟

推理部署时最核心的问题通常是：

• 延迟（latency）：一次请求从输入到输出要多久？
• 吞吐量（throughput/QPS）：每秒可处理多少请求或 batch？

TensorRT 能显著提升这两点，原因是：

• 算子融合减少 kernel 启动和内存访问次数
• 利用 FP16/INT8 让 Tensor Cores 全速工作
• 针对具体硬件调优 kernel 和执行顺序

在多数 CV / NLP 模型上，

• FP16 的 TensorRT 引擎对比原 FP32 PyTorch/ONNXRuntime 通常能有 2--4 倍加速
• INT8 在合适的场景下甚至可以达到 4--8 倍（视模型与 GPU 而定）

2.2 工程：稳定、可控、易于部署

推理阶段的诉求和训练阶段不同：

• 训练： 需求快速迭代、易 debug、多种算子实验等
• 推理： 需求稳定、高性能、可监控、长周期运行

用原生框架直接做推理（例如 PyTorch model.eval()）当然没问题，但在工程上会遇到：

• 框架升级导致行为变化
• 依赖庞大，部署镜像很大
• 没有针对推理做极致优化
• 多语言/多平台支持有限

TensorRT 的特点是：

• 核心是一个 版本相对稳定、ABI 固定的推理引擎
• 可以在 C++、Python 中调用，易于和现有服务框架集成
• 在推理侧只依赖 CUDA/TensorRT，不必把训练框架打进生产镜像

2.3 资源成本：节省显卡和功耗

当模型性能翻倍时，你可以：

• 在同样的 GPU 上支撑 更多用户
• 在同样的负载下，使用更少的 GPU，降低成本

在云上或大规模部署场景里，这一点非常现实。

---

3. TensorRT 的典型工作流程

从训练好的模型到 TensorRT 推理，大致步骤如下：

1. 在训练框架中导出模型为中间格式（通常是 ONNX）
   • PyTorch：torch.onnx.export
   • TensorFlow：SavedModel → ONNX
2. 使用 TensorRT 将 ONNX 构建为 Engine（序列化引擎文件 .plan）
   • 编程方式：用 C++/Python API 的 IBuilder & INetworkDefinition
   • 命令方式：用 trtexec 工具直接：
     trtexec --onnx=model.onnx --saveEngine=model.plan --fp16
3. 在推理服务里加载 Engine 并执行推理
   • 预先分配 GPU 内存
   • 创建 IExecutionContext
   • 通过 enqueueV2() / executeV2() 执行
4. 监控性能与稳定性，必要时重新构建 Engine
   • 换 GPU / 换 TensorRT 版本 / 换 batch size 通常需要重新 build

---

4. 入门学习路线（适合有一定代码能力）

按优先级和难度逐步来：

4.1 第一阶段：快速跑通

目标：从一个 PyTorch 模型 → ONNX → TensorRT → 跑通推理

建议步骤：

1. 选一个简单的模型，例如：
   • ResNet18（图像分类）
   • YOLOv5s（目标检测）
2. 使用 PyTorch 导出成 ONNX
3. 用 trtexec 把 ONNX 转成 TensorRT Engine
4. 写一段 Python 代码：
   • 加载 engine
   • 做一次推理
   • 对比 TensorRT 和原 PyTorch 在同一张图像上的输出

4.2 第二阶段：性能对比与 Profiling

目标：理解 "快在哪里"

• 写脚本对比：
  • PyTorch（CUDA） vs ONNX Runtime（CUDA） vs TensorRT
  • 测以下指标：
    • 单次推理耗时（平均、P95）
    • 不同 batch size 下的吞吐量
• 使用工具：
  • trtexec 内置 profiling 信息
  • Nsight Systems / Nsight Compute
  • TensorRT 的 profiler 回调接口

4.3 第三阶段：精度压缩与 INT8

目标：在可接受精度下降的前提下压榨性能

• 尝试构建 FP16 engine
• 在此基础上再构建 INT8 engine：
  • 需要准备校准数据（calibration dataset）
  • 实现一个量化校准回调，或者使用 PTQ/QAT pipeline
• 对比：FP32 / FP16 / INT8 的：
  • 吞吐、延迟
  • Top-1/Top-5 精度、mAP 等指标

---

5. 实战示例：ONNX → TensorRT → 推理（Python）

下面给一套可执行的最小示例（省略部分错误检查，方便理解结构）。

假定环境：

• 已安装 PyTorch、TensorRT、CUDA
• 有一张 NVIDIA GPU

5.1 用 PyTorch 导出 ResNet18 到 ONNX

# export_resnet18_onnx.py
import torch
import torchvision.models as models

def export_resnet18_onnx(onnx_path="resnet18.onnx"):
    model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
    model.eval().cuda()

    dummy_input = torch.randn(1, 3, 224, 224, device="cuda")
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        export_params=True,
        opset_version=13,
        do_constant_folding=True,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={
            "input": {0: "batch_size"},
            "output": {0: "batch_size"},
        },
    )
    print(f"Exported ONNX to {onnx_path}")

if __name__ == "__main__":
    export_resnet18_onnx()

运行：

python export_resnet18_onnx.py

5.2 使用 trtexec 构建 TensorRT Engine

trtexec \
  --onnx=resnet18.onnx \
  --saveEngine=resnet18_fp16.plan \
  --fp16 \
  --workspace=4096 \
  --minShapes=input:1x3x224x224 \
  --optShapes=input:8x3x224x224 \
  --maxShapes=input:32x3x224x224 \
  --verbose

其中：

• --fp16：启用 FP16
• --min/opt/maxShapes：开启动态 batch size 支持

resnet18_fp16.plan 就是序列化好的 TensorRT 引擎。

5.3 用 Python 加载 Engine 并推理

# infer_trt_resnet18.py
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit  # noqa: F401
import numpy as np
import cv2
import time

TRT_LOGGER = trt.Logger(trt.Logger.INFO)

def load_engine(engine_path):
    with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
        engine = runtime.deserialize_cuda_engine(f.read())
    return engine

def allocate_buffers(engine, batch_size):
    bindings = []
    host_inputs = []
    device_inputs = []
    host_outputs = []
    device_outputs = []

    for binding in engine:
        binding_idx = engine.get_binding_index(binding)
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        shape = list(engine.get_binding_shape(binding))
        if shape[0] == -1:
            shape[0] = batch_size
        size = np.prod(shape)

        # host (page-locked) memory
        host_mem = cuda.pagelocked_empty(size, dtype=dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)

        bindings.append(int(device_mem))
        if engine.binding_is_input(binding):
            host_inputs.append(host_mem)
            device_inputs.append(device_mem)
        else:
            host_outputs.append(host_mem)
            device_outputs.append(device_mem)

    return bindings, host_inputs, device_inputs, host_outputs, device_outputs

def preprocess_image(img_path, input_shape=(3, 224, 224)):
    img = cv2.imread(img_path)
    img = cv2.resize(img, (input_shape[2], input_shape[1]))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.astype(np.float32) / 255.0

    # 简单归一化，可按 ImageNet 规范加减均值/方差
    img = np.transpose(img, (2, 0, 1))  # HWC -> CHW
    return img

def infer(engine_path, img_path, batch_size=1, warmup=10, iters=100):
    engine = load_engine(engine_path)
    context = engine.create_execution_context()

    bindings, host_inputs, device_inputs, host_outputs, device_outputs = \
        allocate_buffers(engine, batch_size)

    # 准备输入 batch
    img = preprocess_image(img_path)
    imgs = np.stack([img] * batch_size, axis=0)  # (B, C, H, W)
    imgs_flat = imgs.ravel()

    # 将输入复制到 host buffer
    np.copyto(host_inputs[0], imgs_flat)

    # 创建 stream
    stream = cuda.Stream()

    def run_once():
        # host -> device
        cuda.memcpy_htod_async(device_inputs[0], host_inputs[0], stream)
        # 推理
        context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
        # device -> host
        cuda.memcpy_dtoh_async(host_outputs[0], device_outputs[0], stream)
        stream.synchronize()

        # 输出 reshape 回 (B, num_classes)
        output_shape = (batch_size, -1)
        outputs = np.array(host_outputs[0]).reshape(output_shape)
        return outputs

    # warmup
    for _ in range(warmup):
        _ = run_once()

    # benchmark
    t0 = time.time()
    for _ in range(iters):
        _ = run_once()
    t1 = time.time()

    avg_ms = (t1 - t0) * 1000 / iters
    print(f"Batch={batch_size}, avg latency={avg_ms:.3f} ms")

    outputs = run_once()
    print("Sample output logits (first sample, first 5 dims):")
    print(outputs[0][:5])

if __name__ == "__main__":
    infer("resnet18_fp16.plan", "test.jpg", batch_size=8)

5.4 对比 PyTorch 推理性能

写一个简单对比脚本：

# benchmark_pytorch_resnet18.py
import torch
import torchvision.models as models
import time

def benchmark_pytorch(batch_size=8, warmup=10, iters=100):
    model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
    model.eval().cuda()
    dummy = torch.randn(batch_size, 3, 224, 224, device="cuda")

    # warmup
    with torch.no_grad():
        for _ in range(warmup):
            _ = model(dummy)

    # benchmark
    torch.cuda.synchronize()
    t0 = time.time()
    with torch.no_grad():
        for _ in range(iters):
            _ = model(dummy)
    torch.cuda.synchronize()
    t1 = time.time()

    avg_ms = (t1 - t0) * 1000 / iters
    print(f"[PyTorch] Batch={batch_size}, avg latency={avg_ms:.3f} ms")

if __name__ == "__main__":
    benchmark_pytorch(batch_size=8)

你可以在同一台机器上分别跑：

python benchmark_pytorch_resnet18.py
python infer_trt_resnet18.py

观察在相同 batch size 下的延迟对比，一般能看到明显的加速（特别是 FP16）。

---

6. 后续进阶 & 验证方向

在你已经能跑通并做简单对比后，可以进一步玩一些更有挑战性的东西：

1. 动态输入尺寸（例如可变分辨率的检测模型）

   • 使用 TensorRT 的 optimization profile
   • 不同分辨率下的性能 & 内存占用对比

2. INT8 量化

   • 使用 TensorRT 的 INT8 calibrator
   • 准备一个小的校准集（几百到几千张图）
   • 对比 INT8 与 FP16/FP32 的：
     • 推理速度
     • 精度（Top1/mAP/IoU）

3. 插件（Plugin）开发

   • 当模型中有 TensorRT 不支持的自定义算子时
   • 用 CUDA 写一个 plugin kernel，将其注册到 TensorRT
   • 流程：实现 IPluginV2DynamicExt → 编译成 .so → 构建引擎时加载

4. 服务化部署

   • 用 Triton Inference Server + TensorRT
   • 或者用你熟悉的 Web 框架（FastAPI / gRPC / C++ 服务）集成
   • 做到：加载 engine → 接收 HTTP/gRPC 请求 → 返回推理结果

5. Profiling & 性能优化套路

   • 使用 TensorRT profiling 接口，查看每个层的耗时
   • 定位瓶颈（可能是前处理、后处理、数据搬运）
   • 探索：
     • batch size 对延迟与吞吐的影响
     • 流并发（multiple streams）
     • 多 GPU / 多实例并行