爱奇艺在DCN、EDVR等 4K 超分模型上的 TensorRT 10倍加速实践

编者荐语：
为了进一步提高模型推理性能，爱奇艺对TensorRT底层机制做了详细的解析。如何对复杂模型推理进行TensorRT的格式转换。以及TensorRT的int8量化推理内部机制，以及如何更好提升视频推理中int8量化模型的推理精度。
以下文章来源于爱奇艺技术产品团队，作者深度学习平台团队

随着终端播放设备的升级，观众对于视频的品质需求也逐步提升。需求从最开始的高清过渡到4K，最近8K也有开始流行的趋势。除了对于分辨率提升的需求之外，视频在采集的过程中，也难免引入一些瑕疵，如物体运动过快导致的模糊，压缩算法导致的画质降低，拍摄/灯光等参数设置不佳导致的细节缺失，噪点增加等。

经典插帧算法一般采用插值等算法，虽然速度很快，但细节丰富的图片放大之后都会比较模糊，去噪更是困难。深度学习方法的引入，因为其庞大的参数空间，很好的拟合了画质的降噪过程，从而在提升分辨率的时候可以提供更多的细节，实现画质和分辨率的双重提升。

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

01 复杂模型的部署挑战

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

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

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

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

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

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

如何对复杂模型推理进行TensorRT的格式转换。
TensorRT的int8量化推理内部机制，以及如何更好的提升视频推理中int8量化模型的推理精度。

02 复杂模型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和pixel shuffle无法正常转换onnx以及TensorRT，故将对应模块排除在外。将对于指定模块使用torch的nn.moduel重新定义一个新的class，即可将对应的分块用onnx给export出来。

2.2 算子融《合》

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

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

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

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

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

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

03 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运算过程中并不都是整形进行计算，中间是穿插着浮点运算的。

TensorRT在做模型转换的时候，会将权值weight如图3所示，以一个合适的缩放因子SW将原始浮点的范围缩放到[-128,127]的int8整形范围。同时在对模型进行finetune/calibration的时候，会对特定卷积层的输入输出总结出其对应的数据范围，通过Scaled Input /Scaled Output从而能够将原本也是浮点的输入和输出缩放到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。

我们在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所示。

04 性能提升结果

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