10 倍加速！爱奇艺超分辨模型加速实践

关注

发布于: 14 分钟前

随着终端播放设备的升级，观众对于视频的品质需求也逐步提升。需求从最开始的高清过渡到 4K，最近 8K 也有开始流行的趋势。除了对于分辨率提升的需求之外，视频在采集的过程中，也难免引入一些瑕疵，如物体运动过快导致的模糊，压缩算法导致的画质降低，拍摄/灯光等参数设置不佳导致的细节缺失，噪点增加等。经典插帧算法一般采用插值等算法，虽然速度很快，但细节丰富的图片放大之后都会比较模糊，去噪更是困难。深度学习方法的引入，因为其庞大的参数空间，很好的拟合了画质的降噪过程，从而在提升分辨率的时候可以提供更多的细节，实现画质和分辨率的双重提升。

但深度学习模型，相比传统方法，其运行时间大幅的提升，单个视频的处理可能要达到数小时或者数天，难以满足对海量视频进行生产的需求。本文在这种背景下，介绍了爱奇艺在视频 4K 超分模型上进行的优化加速和生产落地实践，将 4K 超分模型的性能在 GPU 上提升了 10 倍。

一、复杂模型的部署挑战

1.1 Nvidia 在模型加速上提供的方法

Nvidia 在 Volta 架构后，引入了 tensorcore，一种 domain specific accelerator [DSA]，用于对深度学习中常见的矩阵运算做加速处理。该加速器取得了极大的成功，大幅的降低了大型深度学习模型的推理时间。但 tensor core 作为 DSA 也继承了其普遍缺点，它的编程逻辑特别复杂，普通的程序员在用底层暴露 API 的方式几乎难以使得 tensor core 达到它的理论运行速度。为了降低编程的门槛，Nvidia 构建了 TensorRT[1]框架，将若干针对特定 input/output tensor shape，用手工打磨汇编的方式，得到对应最优 tensorcore 性能的 kernel，并通过抽象的接口暴露给 TensorRT 上层。

TensorRT 在模型编译时，会针对当前模型的状况，依次在各个合适的内核上运行，最终挑选出耗时最小的内核，固定在最终编译的 TensorRT engine 中。Engine 即为最终的部署 binary，TensorRT 推理时会加载 engine，提取对应的内核名称，以及对应的启动参数，按照模型的推理顺序，依次启动内核从而完成推理过程。

虽然 TensorRT 方便了模型借助 tensorcore 得到极大加速，但是由于 tensor core 本身的复杂性，TensorRT 在对外暴露的接口较少，且核心算子实现目前还是闭源，进行模型深度优化时使用方式还是限制较多。举例来说，目前 TensorRT 内部算子都是以 NCHW 进行开发的且仅支持 NCHW 的 tensor 输入，但 tensor core 底层又需要以 NHWC 进行输入，中间会进行多次 tensor reshape 而降低效率。

1.2 爱奇艺在复杂视频推理模型优化上的实践

为了进一步提高模型推理性能，爱奇艺对 TensorRT 底层机制做了详细的解析。通过本文，您将得到如下的知识点：

a. 如何对复杂模型推理进行 TensorRT 的格式转换。

b. TensorRT 的 int8 量化推理内部机制，以及如何更好的提升视频推理中 int8 量化模型的推理精度。

二、复杂模型 TensorRT 的转换方法

对于 TensorRT 模型部署来说，相信用过的人都碰到一个很头疼的问题，即某个算子不支持，或者对于 torch 模型来说，很多算子是需要开发者使用 CUDA 来自定义实现的。其实这个不是 TensorRT 一家的问题，对于 TensorRT 立志成为的通用深度学习编译器来说，深度学习模型和框架迭代非常快速，各种模型的计算需求层出不穷，想要归一化成为一个通用的 IR 表示是非常的困难，更不用说将模型推向性能的极致。

针对不支持 op 或者自定义 CUDA kernel 的处理方法，我们实践中通常用二字口诀来解决，《拆》，《合》。

2.1 模型《拆》解

就一般意义来说，模型的推理过程，其实就是一个完整计算图的重放过程。既然是计算图，我们将其拆解为子图，并桥接对应的输入输出，那么对于其计算结果来说，应该是没有影响的。

所以在这里，我们使用了一个技巧，将可以 export 为 onnx，并正常转换为 TensorRT 的子图独立出来，编译为 TensorRT 的 engine。然后在一个独立的执行文件中，将子图的 engine 依次 replay，中间原始模型未转换的 op，我们用 CUDA kernel 来进行桥接。

图 1 为对于 EDVR[2]具体的拆解部分。对于 EDVR 来说，主要是 DCN 自定义 op 和 pixelshuffle 无法正常转换 onnx 以及 TensorRT，故将对应模块排除在外。将对于指定模块使用 torch 的 nn.moduel 重新定义一个新的 class，即可将对应的分块用 onnx 给 export 出来。

图 1 EDVR 中对于原始计算图的分割

2.2 算子融《合》

以 EDVR 模型为例，其中有一个无法转换的 op 是 Pixelshuffle。该 OP 是超分网络中一种常见的 upsample 操作，用于在网络的末尾处将相关 feature 的 size 给提升到目标大小。因为无法直接转换到 TensorRT 的 engine 中，所以需要将其独立出来实现为自定义的 CUDA kernel，作为前后卷积部分的桥接单元。

但直接这样操作对加速是不友好的。在前面提到过，TensorRT 中 tensor core 的输入是严格 NHWC 的，且 TensorRT 本身是遵守 NCHW。因此 TensorRT 对于单个 engine 来说，输入会有一个 NCHW 至 NHWC 的转换，在一系列操作之后准备输出之时，又会有一个 NHWC 至 NCHW 的转换。这也就意味着我们的桥接 op 方式会触发三个 kernel 操作，而由于超分的像素尺寸特别大，三个 kernel 各自的运行时间也较长。

图 2 pixelshuffle 在 TensorRT 中的融合

图 2 中蓝色虚线框图部分即为 pixelshuffle 计算部分，中间表示为 TensorRT 的原始的算子转换和 pixelshuffle 桥接的计算图，我们可以看到三个模块间进行了多次 NCHW 和 NHWC 的 reshape，很明显这三个 kernel 其实是可以被合到一起的。因为从卷积 A 的结束到卷积 B 的开始，中间的像素只是按照一定的规律进行了三次重排，中间不涉及到任何的计算，且完全线性。

融《合》原理虽很简单，但实现上却很有技巧。由于 TensorRT 闭源，如何消除重排过程，在不改动 TensorRT 本身的情况下不太可能。因为在推理代码中，前后的连续卷积的桥接其实就只有 pixelshuffle 一个 kernel 而已，推理代码并不知道 TensorRT 内部俩个“冗余”重排内核的存在。

所以在这里，又要进行更为细致的拆解。将 TensorRT 拆解为执行文件可以看到的一个个内核，而不仅仅是一个黑盒的 engine 二进制文件。在这里，我们不详述具体的拆解过程，有兴趣的读者可以自行搜索相关 CUDA hook 方法来得到类似的 CUDA kernel 重放方法。

简单的来说，我们用“录音机”将 TensorRT 的运行轨迹给录制成为了“磁带”。并且将磁带按照歌单的顺序把磁带剪成了磁条，并且可以按照新的顺序来重新播放“音乐”。这样原来被隐藏的重排内核就暴露在了执行文件面前，按照其逻辑，我们手工优化了这一内核，从而代替了三个内核的运行。

三、TensorRT 的 int8 推理

在 Volta 架构刚刚推出 tensor core 加速器的时候，Nvidia 只支持了对 fp16 的支持。对于 int8 的完整支持，是到了后面 Turing 架构开始添加，int4/int1 更是到了 ampere 架构才加入。Tensor core 虽然是 Nvidia 应对一众深度学习硬件加速器成功的反击，但软件支持如前所述，有着很大的使用限制。在量化方面，同样如此。

TensorRT 对于量化的支持要比其对 tensor core 的支持更早。大概从 TensorRT4 开始，就已经开始了对 int8 量化的支持[3]。最开始的量化是用从 Pascal 架构开始支持的 dp4a 指令来实现的，该指令可以将 32bit 运算并行化为 4 个 8bit 运算，从而使得 int8 推理速度在当时架构上得到一个质的飞跃。

在 TensorRT 支持 tensor core 之后，也一直沿用着当时的框架，但只支持后训练量化，采用 KL 散度逼近的方式，从全精度模型求解出对应的 int8 量化模型。这不是说 TensorRT 内部就不支持除后训练以外的其他方法，但对于一般用户来说，所能接触的只有 TensorRT 暴露的 API，API 不支持后训练量化以外的方法，那么就意味着普通用户与其他方法的绝缘。

3.1 TensorRT int8 推理的内部机制

从 TensorRT7 开始，Nvidia 开始将 int8 的量化过程以更精细的方式暴露出来。这个一方面也是由于 torch/tensorflow 开始支持了伪量化过程，另外 onnx 也提供了伪量化的 op 表示形式。但很遗憾的是 TensorRT7 对于这种全新 int8 的转换方式支持还是有问题的，其中一个最大的问题就是卷积中的 bias 系数转换的时候弄错了，本应该乘的系数，变成了除，导致加入 CNN 中如果卷积有 bias，那么它的精度将大幅下降。【注：该问题已经在最新的 TensorRT8 中修复】

由于业务模型上线的压力，不可能等到 Nvidia 出下一个版本来修复这个问题，于是我们又将视角投向了拆解。

为了详细的了解 TensorRT int8 的运算过程，我们对 int8 卷积内核做了反汇编，并对汇编代码做了详细的解读。在了解了对应的汇编代码之后，我们明白了其实 int8 运算过程中并不都是整形进行计算，中间是穿插着浮点运算的。

图 3 量化缩放过程

TensorRT 在做模型转换的时候，会将权值 weight 如图 3 所示，以一个合适的缩放因子 SW 将原始浮点的范围缩放到[-128,127]的 int8 整形范围。同时在对模型进行 finetune/calibration 的时候，会对特定卷积层的输入输出总结出其对应的数据范围，通过 SI/SO 从而能够将原本也是浮点的输入和输出缩放到 int8 范围。在 engine 二进制文件内部，权值 weight 就已经被固化成为了 int8 的形式，并且输入在进入卷积层之前就已经是 int8 的形式。这样在卷积算子中的输入/权值卷积乘加计算中，就都是 int 的形式，而由于 int8 的乘积非常容易产生溢出，卷积的乘加累加过程是以 int32 形式存在的。最终乘加的结果也是 int32.乘加之后需要和 bias 进行进一步加法计算。但之前的 int32 结果其实是有输入和权值对应的缩放因子 SW/SI 在里面，所以为了和 bias 的比例一致，需要将原来乘加结果的 int32，先转换为浮点，然后依次除以 SI 和 SW，再加上 bias 才能得到正确的结果，此时仍为浮点类型，为了以整形进行输出，需要乘以 SO 才得到最终输出结果。

这里以公式来表示即：

(IQ*WQ*SI/SW+B)*SO

该公式进而转换为

IQ*WQ*SO*SI/SW+B*SO

在这里，TensorRT 做了一个优化，它将 SO*SI/SW 合并为一个新的系数，并直接将 B*SO 的结果存储在 engine 二进制文件中，这样卷积乘加后只再需要一次 FMA 操作即可获得最终的结果。整个过程可以参加图 4。

图 4 TensorRT int8 内核内部数值分布情况

我们在 TensorRT7 的基础上，抽取了内核的参数，并给其中的权值/bias 进行重新赋值，来解决了原来实现中的问题。

3.2 进一步提升 int8 推理的精度

Int8 虽然大幅的提升了推理的效率，同时 QAT 量化相比 PTQ 来说提升了量化的精度，但整体上 int8 相比全精度来说，仍然要有所下降。为了进一步的提升 int8 模型的推理精度，我们采用了将 TensorRT 内核嵌入 finetune 过程和实时缩放因子计算这两种方法。

TensorRT 内核嵌入 finetune 过程

QAT finetune 过程是一个伪量化的过程。对于 pytorch 来说，它仅仅是在输入和输出上做了一些事情，将权值/输入/输出按照图 x 的缩放过程计算了一番，并缩放到整形范围后，用 round 进行了一次截断处理，随后又使用相同的系数返回之前的浮点。但这一番操作引入的量化误差就比较好的模拟了实际硬件中部署算子的操作。

虽然 pytorch 本身的量化是尽量的模拟实际量化推理计算的过程，但和实际推理计算相比，例如 TensorRT 的 int8 算子的计算结果还是有差异，这个差异对于最终的推理来说就是一个误差的来源。

为了消除这个误差来源，我们将 TensorRT 内核嵌入 pytorch 量化训练的 finetune 过程中，确保量化训练时使用的计算算子和 TensorRT 推理时的算子一致。简单的说，就是将之前的“录音机”又拿过来了，且因为是训练过程使用，内核的权值参数要计算前根据当前的数值进行更新。

缩放因子实时计算

在 3.1 阐述 TensorRT 内部 int8 计算机制的时候，提到过输入/输出的缩放因子是根据 finetune 数据集得到的一个经验数值范围计算得来的。这个过程存在另外一个误差，即实际推理的时候，很有可能视频帧的内容差异导致卷积生成 feature map 的数值范围和 finetune 数据集中的范围不一致。这样继续沿用老的数值，就会导致新的内容推理结果精度的下降。

为了解决这个问题，我们引入了对于缩放因子，特别是输出缩放因子的动态更新。一般来说，对于输入（即原始帧），它的缩放因子都是可以固定的，而如果输出缩放因子可以动态计算出来，就可以由级联关系继承下去变为下层卷积的输入缩放因子，从而整个网络的各个缩放因子得到更新。

在这里，我们又进一步拆解了 TensorRT，给原始 int8 内核增加了一个新的汇编模块，在这里 int8 卷积的输出不再是 int8，而变为 float16 类型。这样的改变使得对于卷积算子来说就不再需要输出的缩放因子的参与。在紧跟着的内核中，使用 reduce 方式计算出整体 float16 的最大值，从而确定出该输出的缩放因子的数值，进而将该 float16 输出缩放为 int8 供下一层输入使用。整体过程如图 5 所示。

图 5 整体 int8 精度优化过程

四、性能提升结果

在整个 EDVR 部署优化的过程中，除了文章提到的《拆》《合》优化之外，同时也包含了其他的一些优化。如图 6 所示中的，我们在优化的第二个步骤中，集中优化了 DCN 自定义 op 中的冗余显存访问，从而大幅的提升了自定义算子自身的效率。在第三个步骤中，我们将一些算子如 leaky 进行融合，随后获取了大约 150ms 的收益。第四以及第五步骤中，我们集中对于一些中间态的格式转换做了相互消减操作，使得在 fp16 精度上，EDVR 达到了 380ms 的速度，最终 int8 的成功应用使得模型的推理效率进一步提升至 1080p 上 180ms 单帧的速度。