第四章：TensorRT推理优化：模型图优化与层融合，精度校准（FP16/INT8/INT4），动态形状与批量处理，TensorRT Python/C++ API实战

好，咱们进入正题。这一章讲的是TensorRT推理优化，说白了就是让模型跑得更快、更省显存。我最早接触TensorRT是在一个自动驾驶项目里，模型在GPU上跑实时检测，帧率死活上不去。后来发现，不是模型不行，是推理引擎没调好。嗯，这一章就是帮你把这块补上。

4.1 模型图优化与层融合

先说说图优化。TensorRT拿到一个模型后，会先做一件事：把计算图重新梳理一遍。它会合并相邻的卷积、BN、ReLU，把多个小算子拼成一个大的kernel。为什么这么做？因为每次调用kernel都有开销，合并后能减少显存读写和启动延迟。

我个人习惯，在转换模型前先用trtexec工具看一眼原始图结构。你会发现很多冗余节点，比如一些框架自动插入的reshape、transpose。这些在训练时有用，推理时就是累赘。

核心优化点：

  • 垂直融合：Conv + BN + ReLU → CBR（一个kernel搞定）
  • 水平融合：多个相同大小的卷积合并成一个更大的卷积
  • 常量折叠：把只依赖常量的计算提前算好，运行时直接取结果
  • 死节点消除：删除训练时有用、推理时无用的分支

我在项目中遇到过一个问题：一个YOLOv5模型转TensorRT后，推理速度反而变慢了。排查后发现，模型里有个自定义的Focus层，TensorRT不认识，把它拆成了很多小操作。后来我手动写了一个plugin，把Focus层合并成一个kernel，速度直接翻倍。你想想看，有时候框架的自动优化并不完美，需要人工介入。

4.2 精度校准：FP16/INT8/INT4

精度校准是TensorRT最核心的功能之一。说白了，就是用更少的数据位来表示模型参数和激活值，从而减少计算量和显存占用。

4.2.1 FP16推理

FP16是最简单的优化方式。你只需要在构建engine时设置BuilderFlag.FP16，TensorRT会自动把能转成FP16的层都转过去。我建议所有项目都先试FP16，几乎无损，速度提升明显。

4.2.2 INT8量化

INT8就复杂多了。它需要校准数据集，用来统计每层激活值的分布，然后找到最佳的量化参数。我曾经在一个分类模型上做INT8量化，结果精度掉了5个点。后来发现是校准集选得不对——我用的是ImageNet的验证集，但模型实际部署的场景是工业缺陷检测，数据分布完全不同。

避坑指南：校准集必须和实际部署场景的数据分布一致。我曾经因为偷懒用了通用数据集，结果线上模型精度崩了。后来老老实实从生产环境抽了1000张图做校准，问题解决。

INT8量化的关键步骤：

  1. 准备校准数据集（500-1000张足够）
  2. 创建Int8Calibrator，实现get_batch()方法
  3. 设置BuilderFlag.INT8并传入校准器
  4. 构建engine，TensorRT会自动完成校准

4.2.3 INT4量化

INT4是最近才支持的，主要用于边缘设备。说实话，精度损失比较大，我一般只在显存极度受限的场景下用。比如一个OCR模型，在Jetson Nano上跑，FP16显存不够，INT8勉强能跑，INT4才能同时跑两个模型。

精度模式 速度提升 精度损失 适用场景
FP32 1x 无 精度要求极高的场景
FP16 1.5-2x 几乎无 大多数场景，首选
INT8 3-4x 0.5-2% 对精度不敏感的场景
INT4 5-6x 3-5% 边缘设备、显存受限

4.3 动态形状与批量处理

动态形状是TensorRT 8.0以后才稳定支持的功能。以前你必须固定输入尺寸，比如1x3x224x224。现在你可以设置一个范围，比如min=(1,3,128,128)，opt=(4,3,224,224)，max=(8,3,512,512)。

我建议你设置opt为实际部署中最常见的尺寸。因为TensorRT会根据opt尺寸做最优的kernel选择，其他尺寸会fallback到通用kernel。我有个项目，输入图片尺寸变化很大，从320x320到1280x1280都有。我把opt设成了640x640，因为这是最常见的尺寸。结果小尺寸和大尺寸的推理速度都还不错。

小技巧：动态形状的engine构建时间比固定形状长很多。我一般先用固定形状调试，最后再切到动态形状。另外，动态形状不支持所有算子，如果遇到不支持的情况，可以用setDynamicRange()手动指定。

批量处理方面，TensorRT支持动态batch。你可以在execute()时传入不同的batch size，只要在构建时指定的范围内就行。我个人习惯把batch size设成2的幂次，比如1、2、4、8，这样对GPU利用率最友好。

4.4 TensorRT Python/C++ API实战

终于到实战环节了。我分别给你展示Python和C++的典型用法。先说Python，它适合快速原型验证。

4.4.1 Python API示例

import tensorrt as trt
import numpy as np

# 创建logger和builder
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)

# 创建网络定义
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)

# 加载onnx模型
parser = trt.OnnxParser(network, logger)
with open("model.onnx", "rb") as f:
    parser.parse(f.read())

# 配置builder
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
config.set_flag(trt.BuilderFlag.FP16)  # 开启FP16

# 构建engine
serialized_engine = builder.build_serialized_network(network, config)

# 保存engine
with open("model.engine", "wb") as f:
    f.write(serialized_engine)

这段代码我用了很多次，基本是标准模板。注意EXPLICIT_BATCH这个flag，不加的话TensorRT会默认batch是隐式的，很多动态形状功能用不了。

4.4.2 C++ API示例

C++ API和Python类似，但更底层一些。我一般在生产环境用C++，因为启动更快，内存控制更精细。

#include <NvInfer.h>
#include <NvOnnxParser.h>

// 创建builder
nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
const auto explicitBatch = 1U << static_cast<uint32_t>(
    nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(explicitBatch);

// 解析onnx
auto parser = nvonnxparser::createParser(*network, logger);
parser->parseFromFile("model.onnx", 
    static_cast<int>(nvinfer1::ILogger::Severity::kWARNING));

// 配置
nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 30);
config->setFlag(nvinfer1::BuilderFlag::kFP16);

// 构建
nvinfer1::IHostMemory* serializedModel = builder->buildSerializedNetwork(*network, *config);

// 保存
std::ofstream f("model.engine", std::ios::binary);
f.write(reinterpret_cast<const char*>(serializedModel->data()), 
        serializedModel->size());
f.close();

嗯，这里要注意一点：C++ API的内存管理是手动的。你创建的builder、network、config都需要手动destroy()。我刚开始用的时候经常忘记，导致内存泄漏。后来我写了个RAII包装类，自动管理生命周期。

4.4.3 推理执行

构建好engine后，推理就简单了。Python和C++的流程一样：

  1. 反序列化engine
  2. 创建execution context
  3. 分配输入输出buffer
  4. 调用execute_v2()或enqueue_v2()

我个人推荐用enqueue_v2()，因为它支持异步执行，可以和CUDA stream配合，实现流水线并行。我曾经在一个视频流处理项目中，用异步推理把吞吐量提升了30%。

总结一下：

  • 图优化是基础，能自动做的让TensorRT做，不能做的自己写plugin
  • 精度校准从FP16开始，INT8需要校准集，INT4慎用
  • 动态形状设置opt为最常见尺寸，batch用2的幂次
  • Python适合原型，C++适合生产，异步推理是性能关键

好了，这一章的内容就这些。下一章我们会讲TensorRT的高级特性，比如多流推理、DLA加速、以及自定义plugin的编写。到时候我会分享一个我踩过的坑——自定义plugin的shape推导写错了，导致整个engine构建失败，排查了两天才找到原因。嗯，到时候细说。

← 上一章 📖 返回目录 下一章 →