三. TensorRT基础入门-TensorRT内部的优化模块

前言

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

本次课程我们来学习课程第三章—TensorRT 基础入门，一起来学习 TensorRT 内部优化的策略

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

在这里插入图片描述

0. 简述

本小节目标：简单理解 Tensor 中每个优化策略，为什么可以做层融合以及理解量化的意义

这节我们学习第三章节第三小节—TensorRT 内部的优化模块，这个小节跟大家介绍一下 TensorRT 内部的优化都有哪些以及内部优化策略比如层融合、量化、调度等等是怎么做的

1.TensorRT的优化策略

首先我们来回顾一下 TensorRT 内部的优化策略，如下图所示：

在这里插入图片描述

TensorRT的各类优化策略

2. Layer Fusion

我们一个个看，我们先看层融合，层融合主要包括以下两种方式：

Vertical layer fusion（垂直层融合）
- 用的比较常见，针对 Conv+BN+ReLU 进行融合
Horizontal layer fusion（水平层融合）
- 当模型中有水平方向上比较多的同类 layer 会直接进行融合

在这里插入图片描述

上图中有三个 kernel，分别是 kernel1、kernel2、kernel3，在 kernel 之间有一个 kernel launch 的空档期，存在一个延迟，所以我们就希望减少这个延迟，那如果这三个 kernel 访问的是同一块内存的同一片区域，那么我们就想能不能只访问一次数据，一次就把计算全部做了呢？一个非常直观的方法就是把 kernel1、kernel2 以及 kernel3 全部都放到一个 kernel 上进行实现

所以层融合的优点就在于它可以减少启动 kernel 的开销与 memory 操作，从而提高效率，同时有些计算可以通过层融合优化后跟其它计算合并

在这里插入图片描述

TensorRT vertical layer fusion

在这里插入图片描述

上图是 TensorRT 中垂直层融合的一个典型案例，将 Conv、BN、ReLU 进行层融合

在这里插入图片描述

TensorRT horizontal layer fusion

在这里插入图片描述

上图是 TensorRT 中水平层融合的一个典型案例，将水平方向上的 3 个 1x1 conv 合并，水平层融合在 transformer 和 self-attention 中还是比较常见的，在 transformer 中 q、k、v 的生成过程会用到水平的层融合

我们思考下为什么 Conv、BN、ReLU 三个算子可以进行垂直层融合呢？这其实是与 BN 层的计算有关的，我们先回顾一下 Batch Normalization 的公式，如下所示：

在这里插入图片描述

其中：

$\mu_B$ ：代表均值
$\sigma_B^2$ ：代表方差
$\epsilon$ ：一个大于 0 的浮点数，用于防止分母为 0
$\gamma$ ：scaling，缩放因子
$\beta$ ：shift，偏移因子

BN 层就是把一个 batch 的图片减均值除以方差，另外需要缩放因子 $\gamma$ 和偏移因子 $\beta$ 参与训练学习的过程，而上面提到的这些数值在推理的时候都是可以提前计算出来

接下来我们把 Conv 和 BN 结合起来看一下，Conv 过程就是一个线性变化，融合 BN 后的整个推导过程如下所示：

在这里插入图片描述

接着再进行展开如下所示：

在这里插入图片描述

其中 $\hat{w}$ 和 $\hat{b}$ 这两个参数值可以提前计算出来，所以你的 convolution 和 batch normalization 融合之后依旧是一个线性变换，只是线性变换的参数发生了相应的变化

那我们接着看 ReLU 为什么可以和 Conv+BN 融合呢？

在这里插入图片描述

ReLU function

ReLU 的函数如上图所示，我们可以看到 ReLU 只是做一个截取大于 0 的值，内部没有计算，因此它也可以融合到 Conv+BN 的线性变换公式里面

Tips：最近的很多模型经常会有很多种类的 activation function，比如 GELU、Swish、Mish 等等，这些激活函数往往由于计算复杂很难加速，可以尝试改成 ReLU 看看精度和改变后性能的提升

3. Kernel Auto-Tuning

我们要讲的另外一个 TensorRT 优化策略是 Kernel auto-tuning，前面我们利用 CUDA 做矩阵乘法计算 $M\times N=P$ 的过程中需要手动去调整 grid size、block size 等大小来得到最佳的性能

在这里插入图片描述

GPU的GEMM计算

但实际上 TensorRT 内部对于同一个层会自动使用不同 kernel 函数进行性能测试

比如对于 FC 层中的矩阵乘法，根据 tile size 有很多种 kernel function
例如 32x32、32x64、64x64、64x128、128x128，针对不同硬件 TensorRT 有不同策略

在这里插入图片描述

Volta architecture block diagram

Tips：有时你的模型架构由于计算密度不够高的原因，在 kernel auto tuning 时 TensorRT 认为没必要选择 Tensor core 而给你分配 CUDA core 使用

4. Quantization

我们再来看 TensorRT 中非常重要的一个优化策略—量化，这里我们简单介绍下大家有个基本印象就像后续课程我们会详细来介绍。量化是压缩模型的一个很重要的策略，它主要是将单精度类型（FP32）训练权重转变为半精度（FP16）或者整型（INT8，INT4），如下图所示：

在这里插入图片描述

FP32->INT8的量化

我们如何表示浮点数呢？例如圆周率 3.1415926… 该如何表示呢？很显然我们用 8bit 是无法表达这种小数的，IEEE 745 standard 规定了一种表示浮点数的方法，它将浮点数分为三个部分：

sign
expond（range）
significant（fraction，precision）

在这里插入图片描述

FP16表示方法

对于浮点数 0.00368 我们该如何用 IEEE 745 standard 来表示呢？其中的 sign、expond、significant 分别是多少呢？我们如果用 FP16 来表示 0.00368，它其实是 $-1^0\cdot3.68\cdot10^{-3}$ ，因此 FP16 中的 sign=0、expond=-3、significant=3.68

我们来具体看一个例子，浮点数 −6.650390625… 如果用 FP16 来表示，那么它的 sign、expond、significant 三部分内容分别是多少呢？具体表示方法如下图所示，大家可以自己推导下：

在这里插入图片描述

那表示浮点数的方式有很多，上面提到的 FP16 只是其中的一种，那还有 FP32、TF32、BF16 等等，如果浮点数所占用的 bits 越多则表示的精度越高，但随之而来的就是占用的内存空间越大，计算时间越长

在这里插入图片描述

训练的时候因为需要优先考虑精度而不太需要重视速度，所以会使用 FP32 来表示权重和激活值，但在部署的时候，我们需要想办法把 FP32 的数据尽量压缩，因此会用16bits、8bits 甚至 4 bits 来表示它们，这个过程就被称作量化（Quantization）

在这里插入图片描述

那其实在 Training 的时候为了加快速度也有使用两个的技巧，这个方法被称作 Mix precision training（混合精度学习），大家感兴趣的可以看看 Baidu 和 NVIDIA 发表在 ICLR2018 上的 Mixed precision training 这篇 paper

Note：由于量化太过于重要，技巧很多，并且也是目前模型部署最优先考虑的策略，所以下一章节会单独拿出来详细介绍，目前只需要知道量化是干什么的就 OK 了。

总结

本次课程我们主要学习了 TensorRT 的优化策略，并对其中的层融合、Kernel 调度以及量化三个优化策略进行了简单的介绍。层融合的优点就在于它可以减少启动 kernel 的开销与 memory 操作从而提高效率；TensorRT 内部还会对于同一个层自动使用不同 kernel 函数进行性能测试以此来选择适合当前硬件的最佳配置；量化则是一个非常重要的策略，它可以压缩模型同时提高计算速度

OK，以上就是第 3 小节有关 TensorRT 内部的优化模块的全部内容了，下节我们来学习剖析 ONNX 架构，敬请期待😄