TensorRT 推理加速笔记

最近在看模型部署这块，绕不开 TensorRT。它和训练框架不是一个层面的东西，定位很明确：对一个已经训练好的神经网络做 inference 的加速。把它的思路理一理。

它在优化什么

TensorRT 的加速主要靠两件事——合并层（combining layers）和自动选择最优 kernel（optimizing kernel selection）。优化的目标可以拆成这么几个推理时真正关心的指标：

  • 延迟（latency）
  • 吞吐量（throughput）
  • 能效（efficiency / 能耗）
  • 内存占用（memory usage）
  • 精度（accuracy）

在硬件允许的情况下，它还可以混合比特、用低精度（比如 FP16/INT8）来进一步压时间和显存。

整体框架大致如下，中间这几块就是主要的优化点：

TensorRT 框架与优化点

graph 上的几个优化点

TensorRT 拿到网络的计算图之后，会在 graph 这一层做一系列改写：

  • 消除无用层：输出根本没被用到的层直接删掉。
  • 层融合：把 convolution、bias、ReLU 这种常见的连续操作融合成一个，省掉中间结果的读写。
  • 同源聚合：把参数足够相似、且来自同一个源张量的操作聚到一起算——典型的就是 GoogLeNet/Inception 模块里那些并排的 1×1 卷积。
  • 合并 concat 层：不真的去做拼接，而是直接把上游各层的输出写到最终该去的地方，省掉一次显式 concatenation。

大致的使用流程

TensorRT 可以把各种框架（TensorFlow、Caffe、ONNX 等）构建的模型转换成它自己的 engine。用起来 C++ 和 Python 都行，主要就这么几步：

  1. 创建一个 TensorRT network definition（描述网络结构）
  2. 调用 TensorRT 的 builder，把 network 编译成 engine
  3. 对 engine 做序列化 / 反序列化（这样不用每次都重新 build，部署时直接加载）
  4. 给 engine 喂数据，执行 inference

简单说，TensorRT 干的活就是：拿走训练好的模型，在图层面把能省的都省掉，再针对当前 GPU 选一套最快的 kernel，必要时降精度，最后打包成一个可以直接部署、反复加载的 engine。

  • TensorRT
  • 推理加速
  • 模型部署
  • 深度学习