前言

自动驾驶之心推出的 《CUDA与TensorRT部署实战课程》，链接。记录下个人学习笔记，仅供自己参考

本次课程我们来学习课程第三章—TensorRT 基础入门，一起来了解 TensorRT

课程大纲可以看下面的思维导图

0. 简述

本小节目标：理解 TensorRT 的工作流程，TensorRT 的优化极限，以及其他 DNN 编译器

这节开始我们进入第三章节-TensorRT 基础入门的学习，第三章节内容主要分为以下五个部分：

• TensorRT 简介
• TensorRT 的应用场景
• TensorRT 的模块
• 导出 ONNX 以及修改 ONNX 的方法
• 初步使用 TensorRT

这次课程我们来学习 TensorRT 简介，这次课程结束之后希望大家理解 TensorRT 的工作流程，TensorRT 的一些局限以及与其它 DNN 编译器相比优势和劣势分别是什么

1. 什么是TensorRT

首先什么是 TensorRT，官方定义如下：

NVIDIA® TensorRT™, an SDK for high-performance deep learning inference, includes a deep learning inference optimizer and runtime that delivers low latency and high throughput for inference applications.

我们可以把 TensorRT 理解为一种针对 NVIDIA GPU 的一种优化编译器，它具备以下几个功能：

• 自动优化模型
  • 寻找模型中可以并行处理的地方
  • 针对当前部署的 GPU 框架，寻找最优的调度和并行策略
• 支持多框架输入
  • ONNX
• Python/C++ API 接口
  • 可以方便在自己的程序中调用 TensorRT API 来实现推理

现在的大多数 DL framework 都可以通过 TensorRT 生成推理引擎部署在硬件上

• 比较常见的组合：Pytorch->ONNX->TensorRT
• 目前自动驾驶的部署硬件大多都会采用 NVIDIA
  • DRVIE series
  • Jetson series

很多厂商除了 NVIDIA 也会有很多其它选择，比较典型的是高通的 Heterogeneous Architecture 如下图所示：

Qualcomm

另外一个比较小众的是 Hailo 如下图所示，它是中东那边开发的一个比较典型的小型 Edge Device for DNN，它的功耗低且效率高，有很多量化和调度的技巧，此外它针对自己的硬件而设计的 Dataflow Compiler 也很精妙

Hailo

Note：Dataflow 数据流是一个很传统的概念，最近针对 DNN 的硬件设计有偏向 Dataflow 的趋势。大家感兴趣的同学可以自学一下，建议参考 Hailo, Google TPU, MIT Eyeriss 的 paper

Google 的 TPU 是属于 dataflow architecture 的一个例子，TPU 内部有一个 on-chip memory 叫做 HBM（high-bindwidth memory），当读取数据的时候，HBM 里的数据会以 queue 的形式放到 MXU（Matrix Multiplication Unit）进行数据的流动以及计算

为了理解 TPU 的工作原理，我们不妨了解一下其他加速器如何应对训练机器学习模型的计算挑战，以下内容 copy 自：https://cloud.google.com/tpu/docs/intro-to-tpu?hl=zh-cn

CPU 的工作方式

CPU 是一种基于冯·诺依曼结构的通用处理器。这意味着 CPU 与软件和内存协同工作，如下所示：

CPU 最大的优点是它们的灵活性。您可以在 CPU 上为许多不同类型的应用加载任何类型的软件。例如，您可以使用 CPU 在 PC 上处理文字、控制火箭引擎、执行银行事务，或使用神经网络对图片进行分类。

对于每次计算，CPU 从内存加载值，对值执行计算，然后将结果存储回内存中。与计算速度相比，内存访问速度较慢，并可能会限制 CPU 的总吞吐量。这通常称为冯·诺依曼瓶颈。

对于每次计算，CPU 从内存加载值，对值执行计算，然后将结果存储回内存中。与计算速度相比，内存访问速度较慢，并可能会限制 CPU 的总吞吐量。

GPU 的工作方式

为了提高吞吐量，GPU 在单个处理器中包含数千个算术逻辑单元 (ALU)。现代 GPU 通常包含 2500 - 5000 个 ALU。大量处理器意味着您可以同时执行数千次乘法和加法运算。

这种 GPU 架构非常适合并行处理大量运算（例如神经网络中的矩阵运算）的应用。实际上，在用于深度学习的典型训练工作负载上，GPU 的吞吐量可比 CPU 高出一个数量级。

不过，GPU 仍然是一种通用处理器，必须支持许多不同应用和软件。因此，GPU 与 CPU 存在相同的问题。对于数千个 ALU 中的每一次计算，GPU 都必须访问寄存器或共享内存，以读取运算对象以及存储中间计算结果。

TPU 的工作方式

Google 设计了 Cloud TPU，它们是专门用于神经网络工作负载的矩阵处理器。TPU 不能运行文字处理程序、控制火箭引擎或执行银行交易，但它们可以很快地处理神经网络中使用的大量矩阵运算。

TPU 的主要任务是矩阵处理，这是乘法和累加运算的组合。TPU 包含数千个乘法累加器，这些累加器彼此直接连接以形成大型物理矩阵。这称为脉动阵列架构。在单个处理器上，Cloud TPU v3 包含两个 128 x 128 ALU 的收缩阵列。

TPU 主机将数据流式传输到馈入队列中。TPU 从馈入队列加载数据，并将其存储在 HBM 内存中。计算完成后，TPU 会将结果加载到馈出队列中。然后，TPU 主机从馈出队列读取结果并将其存储在主机的内存中。

为了执行矩阵操作，TPU 将参数从 HBM 内存加载到矩阵乘法单元 (MXU) 中。

然后，TPU 从内存加载数据。每次执行乘法运算时，所得结果都会传递给下一个乘法累加器。输出是数据和参数之间所有乘法结果的总和。在矩阵乘法过程中，不需要访问内存。

以下是三者的对比：

• CPU
  • 需要最高灵活性的快速原型设计
  • 训练时间不长的简单模型
  • 有效批量大小较小的小型模型
  • 包含许多使用 C++ 编写的自定义 TensorFlow 操作的模型
  • 受主机系统可用 I/O 或网络带宽限制的模型
• GPU
  • 具有大量自定义且必须至少部分在 CPU 上运行的 TensorFlow/PyTorch/JAX 操作的模型
  • 具有不能在 Cloud TPU 上使用的 TensorFlow 操作的模型
  • 有效批量大小较大的中到大型模型
• TPU
  • 由矩阵计算主导的模型
  • 在主训练循环内没有自定义 TensorFlow/PyTorch/JAX 操作的模型
  • 需要训练数周或数月的模型
  • 有效批量大小较大的大型模型

OK，我们回归正题，市场上有很多部署硬件但为什么 NVIDIA 和 TensorRT 一直都很火呢？主要有以下几个原因：

• 整体硬件设计和编译技术很成熟
  • BUG 算比较少的那种
  • 对量化的支持比较全面以及完善
  • SDK 很充足（polygraph，onnxsurgeon）
• 可以参考的文档比较多
  • 官方给的资源还是相当丰富的
  • 给的 sample 也很适合学习
• Community 很大
  • 出现问题的时候可以讨论

但是：

• 价格偏贵，当面向量产的时候价格是很需要重视的
• 过大的 TOPS 是否真的就很好？
• TOPS（Tera Operations per second）是衡量一个硬件的计算力的常用指标，之后我们会详细讲，现阶段有个认识就好

2. TensorRT的工作流介绍

TensorRT 的工作流如下图所示：

主要可以分为以下几部分：

• 1. 将训练好的模型转为 TensorRT 可识别的模型
• 2. TensorRT 逐个 Layer 进行分析并尝试各种优化策略（量化、层融合、调度、多流执行、内存复用等等）
• 3. 生成推理引擎，可以从 C++/Python 程序中调用

3. TensorRT的一些限制

但是 TensorRT 也存在一些限制，主要有以下几点：

1. 针对不支持的算子

• 看看不同 TensorRT 版本中是否有做了支持
  • Pytorch：Facebook
  • ONNX：Microsoft
  • TensorRT：NVIDIA
  • 部署时一般是 Pytorch->ONNX->TensorRT，我们可以看到中间跨越了三个不同的厂家，所以不同厂家由于算子版本不同带来的兼容性问题会很头痛
  • 比如 LayerNormalization 算子在 TensorRT-8.6-GA 中就有支持，而在之前的版本没有
• 修改 PyTorch 的算子选择，让它使用一些 TensorRT 支持的算子
• 自己写插件，内部实现自定义算子以及自定义 CUDA 加速核函数
• 不使用 ONNX 自己创建 parser 直接调用 TensorRT API 逐层创建网络
  • 比如 tensorrtx 这个 repo

2. 不同 TensorRT 版本的优化策略是不一样的

• 比如对 Transformer 的优化 TensorRT-7.x 和 TensorRT-8.x 跑出来的性能是不一样的
• 如果我们希望好的优化策略那我们就可以考虑比较新的 TensorRT 版本去做模型部署

3. 有时你预期 TensorRT 的优化和实际的优化是不一样的

• 比如说你期望 TensorRT 使用 Tensor core 但 kernel autotuning 之后 TensorRT 觉得使用 Tensor core 反而会效率降低，结果给你分配 CUDA core 使用。因为内部需要做一些额外的处理
• 比如说 INT8 量化的结果比 FP16 还要慢，这是 TensorRT 内部的一些优化调度所导致的

4. 天生并行性差的 layer，TensorRT 也没有办法

• 1x1 conv 这种 layer 再怎么优化也没有 7x7 conv 并行效果好

重点注意：TensorRT 虽然很方便但不要过分的依赖 TensorRT 给你的结果，你需要结合部署的硬件特性做一些 Benchmark 和 Profiling，学习使用一些 Profile tools 去分析模型的计算瓶颈在哪里，以及去分析你的预测优化策略和实际优化策略产生分歧的原因在哪里。

总结

本次课程我们学习了解了 TensorRT，TensorRT 是 NVIDIA 针对 DNN 推出的一个优化编译器，它可以优化经过训练的深度学习模型以实现高性能推理，它的整个工作流程分为三部分，首先转换为 TensorRT 可识别的模型比如 ONNX，接着 TensorRT 尝试各种策略（算子/层融合、量化等）进行优化，最后生成对应的 engine 利用 API 调用。

当然我们也需要认识到 TensorRT 本身是存在一些限制的，我们不要过分的依赖 TensorRT 的结果，而需要学会结合硬件特性利用一些工具来分析模型的计算瓶颈

OK，以上就是第 1 小节有关 TensorRT 简介的全部内容了，下节我们来学习 TensorRT 的应用场景，敬请期待😄

参考

• Qualcomm Software Architecture
• Detail of Hailo AI Edge Accelerator Emerge
• Hailo-15 quad-core AI Vision processor delivers up to 20 TOPS for Smart Cameras
• Introduction to Cloud TPU
• NVIDIA Deep Learning TensorRT Documentation
• NVIDIA TensorRT