第7章 TensorRT部署

课程目录第7章 TensorRT部署

7、TensorRT部署：TensorRT安装与版本检查、ONNX模型转换示例、INT8量化配置、性能基准测试

说到TensorRT，这玩意儿在Jetson平台上简直就是性能加速的核武器。我最早接触它是在一个边缘检测项目上，模型跑在Jetson Nano上，帧率只有可怜的8帧。后来上了TensorRT，直接飙到30帧——嗯，那种感觉就像给自行车装了个小马达。

说白了，TensorRT就是NVIDIA专门为自家GPU做的深度学习推理优化器。它能做层融合、精度校准、内存复用这些事。你想想看，一个训练好的模型，里面可能有一堆冗余计算，TensorRT就是把这些冗余砍掉，让模型跑得更快。

7.1 TensorRT安装与版本检查

JetPack SDK里其实已经预装了TensorRT。但有个坑——不同版本的JetPack对应的TensorRT版本不一样。我建议你装好系统后，先确认一下版本。

检查版本很简单，打开终端敲一行命令：

dpkg -l | grep tensorrt

你会看到类似这样的输出：

ii  libnvinfer8  8.5.2-1+cuda11.4  arm64  TensorRT runtime libraries
ii  tensorrt     8.5.2-1+cuda11.4  arm64  Meta package for TensorRT

这里显示的是TensorRT 8.5.2版本，对应CUDA 11.4。我个人习惯用python3 -c "import tensorrt; print(tensorrt.__version__)"再确认一下Python接口的版本。

注意： 如果你用的是JetPack 5.x系列，TensorRT版本一般是8.5.x。JetPack 4.x对应的是7.x版本。版本差异会影响ONNX算子支持，后面转换模型时可能会遇到不兼容的问题。

如果需要重新安装，可以用：

sudo apt-get update
sudo apt-get install tensorrt

不过说实话，我建议直接用JetPack自带的版本。自己手动装容易搞乱依赖关系，我曾经在Orin上折腾了一下午，最后发现还是预装的最稳。

7.2 ONNX模型转换示例

ONNX是模型转换的中间格式。你训练好的PyTorch或TensorFlow模型，先转成ONNX，再喂给TensorRT。这个流程我走了无数遍，踩过的坑比走过的路还多。

先看一个PyTorch转ONNX的例子：

import torch
import torchvision.models as models

# 加载预训练模型
model = models.resnet18(pretrained=True)
model.eval()

# 创建示例输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
    opset_version=11
)
print("ONNX模型导出成功")

这里有个关键点——opset_version。我建议用11或13，太高的版本TensorRT可能不支持。动态轴（dynamic_axes）这个参数，如果你需要处理可变batch size，一定要加上。我之前有个项目没加，结果线上推理时batch size一变，模型直接崩了。

ONNX转TensorRT引擎，用trtexec工具最方便：

trtexec --onnx=resnet18.onnx \
        --saveEngine=resnet18.engine \
        --fp16 \
        --workspace=1024

--fp16表示启用半精度推理，--workspace是工作空间大小，单位MB。如果你的模型比较大，可以适当调高workspace。

小技巧： 转换时加上--verbose参数，可以看到每一层的详细信息。如果转换失败，日志里会明确告诉你哪个算子不支持。

7.3 INT8量化配置

INT8量化是TensorRT的杀手锏。把FP32的模型压缩到INT8，推理速度能翻倍，显存占用也大幅降低。但代价是精度会掉一点。

配置INT8量化需要校准数据集。说白了就是给TensorRT一些真实数据，让它统计每层激活值的分布，然后找到最优的量化参数。

看一个Python接口的INT8量化示例：

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

# 创建builder
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, TRT_LOGGER)

# 加载ONNX模型
with open("resnet18.onnx", "rb") as f:
    parser.parse(f.read())

# 配置INT8量化
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

# 设置INT8模式
config.set_flag(trt.BuilderFlag.INT8)

# 设置校准器
class Calibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, batch_size, input_size):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.batch_size = batch_size
        self.input_size = input_size
        self.batch_idx = 0
        self.device_input = cuda.mem_alloc(batch_size * input_size * 4)

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, names):
        if self.batch_idx < 10:  # 使用10个batch做校准
            # 这里应该加载真实数据
            data = np.random.randn(self.batch_size, *self.input_size).astype(np.float32)
            cuda.memcpy_htod(self.device_input, data)
            self.batch_idx += 1
            return [int(self.device_input)]
        else:
            return None

    def read_calibration_cache(self):
        return None

    def write_calibration_cache(self, cache):
        with open("calibration.cache", "wb") as f:
            f.write(cache)

calibrator = Calibrator(1, (3, 224, 224))
config.int8_calibrator = calibrator

# 构建引擎
engine = builder.build_serialized_network(network, config)
with open("resnet18_int8.engine", "wb") as f:
    f.write(engine)
print("INT8引擎构建成功")

  核心要点： 校准数据一定要来自真实场景。我用随机数做过测试，精度掉了5个点。换成真实数据后，只掉了1个点。校准batch数量一般10-20个就够了，太多反而过拟合。

INT8量化后，建议用验证集跑一遍精度对比。如果精度下降太多，可以考虑对某些敏感层回退到FP16。TensorRT支持逐层精度设置，这个后面有机会再细聊。

7.4 性能基准测试

模型部署完，总得跑个分吧？性能基准测试主要看三个指标：延迟、吞吐量、显存占用。

用trtexec做基准测试最省事：

trtexec --loadEngine=resnet18.engine \
        --batch=1 \
        --warmUp=100 \
        --iterations=500 \
        --duration=10

输出结果里重点关注这几项：

指标	说明	参考值（ResNet18, Jetson Orin）
Latency	单次推理延迟	FP32: 8ms, FP16: 4ms, INT8: 2ms
Throughput	每秒处理帧数	FP32: 125 FPS, FP16: 250 FPS, INT8: 500 FPS
GPU Memory	显存占用	FP32: 200MB, FP16: 120MB, INT8: 80MB

我自己写了个简单的Python脚本做基准测试：

import tensorrt as trt
import numpy as np
import time

# 加载引擎
with open("resnet18.engine", "rb") as f:
    engine_data = f.read()
runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
engine = runtime.deserialize_cuda_engine(engine_data)

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

# 分配内存
input_buf = np.random.randn(1, 3, 224, 224).astype(np.float32)
output_buf = np.empty((1, 1000), dtype=np.float32)

# 预热
for _ in range(100):
    context.execute_v2([input_buf, output_buf])

# 正式测试
latencies = []
for _ in range(500):
    start = time.time()
    context.execute_v2([input_buf, output_buf])
    latencies.append((time.time() - start) * 1000)

print(f"平均延迟: {np.mean(latencies):.2f} ms")
print(f"最大延迟: {np.max(latencies):.2f} ms")
print(f"最小延迟: {np.min(latencies):.2f} ms")
print(f"吞吐量: {1000 / np.mean(latencies):.2f} FPS")

避坑指南： 我曾经在测试时没做预热，结果前几次推理延迟特别高，因为GPU还在初始化状态。预热100次左右，等GPU跑稳了再测，数据才真实。

还有一个容易被忽略的点——CPU和GPU之间的数据传输时间。如果你的模型输入很大，数据传输可能成为瓶颈。我建议用cuda.memcpy_htod和cuda.memcpy_dtoh的异步版本，配合CUDA流，能有效隐藏传输延迟。

最后说一句，性能测试一定要在目标平台上跑。在PC上测出1000FPS，放到Jetson上可能只有100FPS。平台差异太大了，别偷懒。

我的经验： 保存每次测试的配置和结果，方便对比。我习惯用CSV记录：日期、模型、精度模式、batch size、延迟、吞吐量。这样调优时能清楚看到每次改动带来的变化。

这张图把整个TensorRT部署流程串起来了。从训练好的模型，到ONNX中间格式，再到TensorRT引擎构建，最后部署推理。INT8量化是中间的一个分支流程，需要校准数据集参与。性能测试则是最后一步，验证你的优化效果。

嗯，TensorRT这部分内容确实不少，但掌握了这些，你在Jetson上的模型部署基本就稳了。记住一句话：先跑通，再优化，最后再上INT8。别一上来就搞量化，容易把自己搞晕。

📚 课程目录

7、TensorRT部署：TensorRT安装与版本检查、ONNX模型转换示例、INT8量化配置、性能基准测试

7.1 TensorRT安装与版本检查

7.2 ONNX模型转换示例

7.3 INT8量化配置

7.4 性能基准测试