探讨TensorRT加速AI模型的简易方案 — 以图像超分为例

LiveVideoStack 2021年4月30日

AI模型近年来被广泛应用于图像、视频处理,并在超分、降噪、插帧等应用中展现了良好的效果。但由于图像AI模型的计算量大,即便部署在GPU上,有时仍达不到理想的运行速度。为此,NVIDIA推出了TensorRT,成倍提高了AI模型的推理效率。本次LiveVideoStack线上分享邀请到了英伟达DevTech团队技术负责人季光一起探讨把模型运行到TensorRT的简易方法,帮助GPU编程的初学者加速自己的AI模型。


文 / 季光

整理 / LiveVideoStack


01 关于NVIDIA GPU


首先介绍英伟达的GPU。上一代GPU架构是图灵Turing,当前架构是安培Ampere。Ampere消费级型号都是30开头,包括3090、3080、3070等;企业级型号用于数据中心,包括A100、A30、A10、A16等。由于企业级型号很多,所以简单介绍一下这些型号的用途。


  • A100是芯片面积最大的GPU,适合做训练;而A30的能力大约是A100的一半。但这两个GPU的特点是它们都支持新的数据格式TF32,并且在Tensor Core上做矩阵乘法有很高的吞吐(见上图表格中标绿处)。TF32在训练时非常有用,可以部分替代FP32。另外A100/A30支持MIG,可在单一操作系统中动态切割成多GPU,也可兼用于推理。

  • A10是T4的替代者,它的特点是FP32/FP16吞吐很高,比较适合做推理。

  • A16比较独特,这个卡上含有4个GPU,每个GPU上带着1个NVENC和2个NVDEC引擎,它更适合做转码。

  • GeForce 3090是消费型号,它的GPU型号与企业级的有所不同,计算能力有所欠缺,例如它的FP16的矩阵乘算力是142 TFLOPS(FP16累加,精度有限)或71 TFLOPS(FP32累加)。相比之下,A10的FP32累加矩阵乘可达125 TFLOPS,比它高出很多 。因此无论是做训练还是做推理,GeForce 3090在很多情况下都比不过企业型号。


02 GPU编程基础


GPU算力的发挥要靠GPU上的程序运行出来,因此需要我们编写GPU的程序。GPU编程又被称作异构编程,与CPU编程有不一样的地方。


对于CPU程序,程序和数据都放在主存(即内存)上,这是我们熟悉的方式。而上图左边则是GPU程序的运行方式。GPU有自己的存储器,即显存。要把程序运行在GPU上时,我们需要先把数据从主存拷到显存上,然后启动GPU程序进行计算;当计算完成时,需要把数据从显存拷回主存。以上就是异构编程的思想。简单来说就是将数据拷至异构的处理器上,启动程序,最后将数据拷回。


上图右边是个比较完整的程序,演示了上述思想。程序用cudaMalloc分配出显存上的变量a和b(由显存指针dp_a和dp_b指向),用cudaMemcpy把a从主存拷贝到显存上,然后启动GPU程序。黄色高亮的这段GPU程序称作CUDA kernel,它所使用的数据都来自显存。计算完成后,cudaMemcpy把结果b拷回主存,最后cudaFree释放起初分配的显存。


掌握“数据拷到显存-启动GPU程序-数据拷回主存”这一思想是非常重要的。对于熟悉C++编程的人来说,调用相关函数比较简单,但要写出CUDA kernel还需要额外花功夫。我们特别希望在使用GPU时可以减轻编程负担,通过API调用方式就让程序在GPU上运行起来。这也是TensorRT这种GPU加速库出现的原因。


03 GPU转码流水线中的TensorRT


前面示例代码中的数据是单个浮点数,这是一种简单场景。而更复杂的场景下,拷贝的数据可以是单张图片或连续图片。无论如何,在主存和显存间拷贝数据是有代价的,在数据量大时会成为程序运行的瓶颈,我们需要尽可能地减少或者避免。


以视频转码为例,如果输入数据是编码过的视频码流,可以利用GPU上的硬件解码器解码,把解出的图片存放在显存,再交给GPU程序处理。此外,GPU上还带有硬件编码器,可以将处理后的图片进行编码,输出视频码流。在上述流程中,无论是解码,还是数据的处理,还是最后的编码,都可以使数据留在显存上,这样可以实现较高的运行效率。


04 用TensorRT加速AI模型推理


深度学习应用的开发分为两个阶段,训练和推理。TensorRT用来加速推理。

TensorRT的加速原理大体在这几个方面:


  • TensorRT可以自动选取最优kernel。同样是矩阵乘法,在不同GPU架构上以及不同矩阵大小,最优的GPU kernel的实现方式不同,TensorRT可以把它优选出来。

  • TensorRT可以做计算图优化,通过kernel融合,减少数据拷贝等手段,生成网络的优化计算图。

  • TensorRT支持fp16/int18,对数据进行精度转换,充分利用硬件的低精度、高通量计算能力。

05 TensorRT的加速效果


我们通过一些例子来说明TensorRT的加速效果。


  • 对于常见的ResNet50来说,运行于T4,fp32精度有1.4倍加速;fp16精度有6.4倍加速。可见fp16很有用,启用fp16相较于fp32有了进一步的4.5倍加速。

  • 对于比较知名的视频超分网络EDVR,运行于T4,fp32精度有1.1倍加速,这不是很明显;但fp16精度有2.7倍加速,启用fp16相较于fp32有了进一步的2.4倍加速。


可以看出不同模型的加速效果不同,一般来说卷积模型加速较为显著,而含大量数据拷贝的模型加速效果一般,且fp16无明显帮助。


06 快速上手TensorRT


TensorRT该怎么用呢?本质上就是把训练框架上训练好的模型迁移到TensorRT上。以下是三种方案:


1)通过框架内部集成TensorRT

TensorFlow集成了TF-TRT,PyTorch还有TRTorch,调用这些API就可以把模型(部分地)运行在TensorRT上。它们的使用方式都比较简单,通过框架中的API就能运行,但是很多情况下没有达到最佳效率。


2)比较硬核的方法是使用TensorRT C++/Python API自行构造网络,用TensorRT的API将框架中的计算图重新搭一遍。这种做法兼容性最强,效率最高,但难度也最高。对于这种方法,我们之前在GTC China做过两次报告(TENSORRT: 加速深度学习推理部署,利用 TENSORRT 自由搭建高性能推理模型
https://on-demand-gtc.gputechconf.com/gtcnew/sessionview.php?sessionName=ch8306-tensorrt%3a+  加速深度学习推理部署
),有兴趣的话可以看一看,其难点是需要了解TensorRT的layer都有哪些,以及从原始框架的OP(即操作)跟这些layer的对应关系。


3)今天推荐的方法是从现有框架导出模型(ONNX)再导入TensorRT。

它的优点是难度适中,效率尚可,可以算作捷径。需要解决的问题是:如何从训练框架导出ONNX,以及如何把ONNX导入TensorRT。


07 解决如何导出与如何导入


第0步:了解TensorRT编程的基本框架


上图展示的代码是TensorRT最基本的使用方法。


1.作为准备工作,先造了logger,又造了builder,从builder造出network,这些对所有TensorRT程序都是固定的。


2.接下来高亮的这一部分是通过TensorRT的API把计算图重建起来,使TensorRT上的计算与训练框架原始模型一模一样。这段代码可以非常长,比如上百行。


3.做完之后利用network可以构建TensorRT engine(build_cuda_engine),构建时间根据网络大小有长有短,短的几秒,长的可达几分钟甚至几小时。


4.构建好engine后可以调用运行。而且engine可以保存到磁盘,在第二次运行的时候,不需要再次build,直接load就可以运行。


上图中的d_input、d_output是前面提到的异构编程中的显存地址。


高亮的这一部分可以非常复杂,但为了省事,我们使用ONNX Parser自动搭建网络,让这一部分自动完成。


所以基本流程是这样:先从训练框架导出ONNX,再用TensorRT自带的工具trtexec把ONNX导入TensorRT构建成engine,最后编写一个简单的小程序加载并运行engine即可。


第1步:从框架中导出ONNX


ONNX是中立计算图表示,PyTorch有TouchScript,TensorFlow有frozen graph,都是框架特有的对于计算图持久化的办法。ONNX是平台中立的,理论上所有框架都可以支持的表示方法。


一般情况下导出的ONNX仍具备运行能力,但有时不能直接运行,而是需要补充ONNX Runtime。比如导出的ONNX中具有特殊的算符,例如Deformable Convolution,它不是ONNX标准OP,但通过扩展ONNX Runtime可以让导出的ONNX跑起来。


但ONNX能不能运行并不是可被TensorRT顺利导入的先决条件。也就是说,导出的ONNX不能跑也没关系,我们仍有办法让TensorRT导入。这一点会在下文举例说明。


上图可以看到PyTorch导出ONNX的示例代码。其中的resnet50是一个PyTorch nn.Module对象;verbose设为True可让ONNX用文本方式打出来,对调试很有用;opset可以设置最高到12,版本越高,支持OP数量越多。


第2步:用Parser将ONNX导入TensorRT


TensorRT官方开发包自带可执行文件trtexec。它可以接受ONNX输入,根据ONNX将TensorRT网络搭建起来,构建engine,并保存成文件。这一系列动作通过图中的命令就可以做到。


其实trtexec也可以自己编程来实现,不过一般来说没有必要。


trtexec运行成功说明TensorRT用自有的层重建了等价于ONNX的计算图,而且计算图被顺利构建成了engine。保存成文件的engine将来可以反复使用。


第3步:运行Engine


最后一个步骤比较简单,就是加载engine文件,提供输入数据,即可运行。C++和Python的示例代码可以从这里找到。(
https://github.com/NVIDIA/trt-samples-for-hackathon-cn)


注意一定要对比TensorRT与原框架的计算结果,算出两者的相对误差均值。理想情况下fp32的误差在1e-6数量级,fp16的误差在1e-3数量级。


另外,我们都很关心模型跑到TensorRT上有多少加速比。熟悉CUDA编程的朋友可以用CUDA event测量运行时间,但要注意stream要设置正确。另外还有一种较粗略的简易方法:做一次GPU同步,然后取时间t0;启动GPU程序;再做一次GPU同步,取t1,得t1-t0,这就是GPU程序的运行时间。

(示例代码见这里:https://github.com/NVIDIA/trt-samples-for-hackathon-cn/blob/master/python/app_onnx_resnet50.py  )


这里关键需要理解GPU同步的含义:GPU程序是从CPU启动的,即在CPU端调用TensorRT的execute函数,其实是把GPU程序放进任务队列,放好了就返回了,并不等GPU程序执行完毕;而GPU程序的执行却是异步的。在CPU上做一次GPU同步,就是让CPU等待此前提交的GPU任务全部执行完。基于以上,我们就可以理解为什么取时间之前要做一次GPU同步。


这里有个问题:这个简易方法在什么时候不准确?简单的说,这个方法会有误差,如果要统计的GPU程序运行时间较短,就很难得出准确结果。这种时候,用CUDA event才是终极解决方法。


08 导出ONNX:疑难问题


前面说得都是最顺利的情况。我们看看对于导出ONNX,不顺利的情况有哪些:


如果遇到ONNX不支持的操作,解决办法是升级框架和ONNX导出工具,使用当前支持的最高opset。


但这样可能还不够,因为有些PyTorch官方的OP在ONNX中仍然没有定义(或无法组合得到)。所以在导出时加上选项ONNX_FALLTHROUGH,即便没有定义也可以导出。


如果遇到开发者自定义的OP,则需要确认为自定义的Function子类增加symbolic函数,从而为自定义OP取ONNX节点名。

(例子见这里:https://github.com/shining365/EDVR-TRT/blob/master/basicsr/models/ops/dcn/deform_conv.py#L114  )


此外,用trtexec把ONNX导入TensorRT时可能会遇到报错。一种常见的情况是不支持的OP,这个稍后再说。另一种情况是TensorRT Parser对ONNX网络结构有特殊要求。具体地,我们看一个例子。


上图中高亮的报错信息是”Resize scales must be an initializer!”为了得到更丰富的信息方便调试,请运行trtexec时打开--verbose选项。从图中可以看到,这个Resize节点有385,402,401这3个输入。这3个数字并不是输入的具体值,而是输入变量的名字。我们需要进一步看看这3个变量都是怎么生成的。


请在导出ONNX时确保设置verbose=True,可得到文本描述的ONNX,见上图。可以看到Resize节点在图中最下方,它的3个输入变量已被高亮出来,它们有各自的计算过程。由于ONNX本身是个计算图,我们可以画一张图将这一部分更清楚地展现出来。


09 ONNX手术刀:Graph Surgeon


上图是有关这个Resize的ONNX子图。它的第三个参数变量401来自Concat操作,将3个变量Concat在一起:其中一个是Constant,另外两个是Constant经过了Unsqueeze与Cast,做了数据类型的转换。


前文报错信息“Resize scales must be an initializer!”指的是Resize的第三个参数不能是变量,而必须是Constant,所以我们需要把蓝色的这部分子图转换成一个Constant,变成右边的样子。一旦做到,TensorRT Parser就会正常运行下去。


这个转换在理论上可以做到,原因是这部分子图的叶子节点都是Constant,具体值都写在里面,我们按计算图手工做一下相关计算,得到结果后存放在新建的Constant节点里就可以了。实现它的工具是Graph Surgeon。


Graph Surgeon像手术刀一样可以修改ONNX计算图。上图就是用Graph Surgeon完成计算图转换的代码。


1.首先找到符合条件的Resize节点,其筛选条件就是它的第三个输入变量应来自Concat节点。


2.然后我们对这个Concat的所有输入参数建立一个while循环,一直往上走,直到找到Constant,并把Constant里面的值放进values中。这样走完for循环后,所有要合并的值都已经存进values中。


3.最后新建Constant节点,用numpy的concatenate函数将值合并填入该节点,并为该节点连接好输出。


Graph Surgeon的一个完整示例代码见这里(
https://github.com/NVIDIA/trt-samples-for-hackathon-cn/blob/master/python/app_onnx_custom.py  )
。随着大家做Graph Surgeon的经验积累,特殊情况处理的经验会越来越丰富,你将会积累更多的节点处理方式,从而让更多模型被TensorRT Parser正确解析。


10 遇到不支持的操作


当trtexec报告不支持的OP时,我们不得不编写TensorRT Plugin。TensorRT Plugin是TensorRT功能的扩展,需要什么我们就可以写什么,也可以说是“万金油”。


编写TensorRT Plugin的思想是套用模板在里面“填空”。最关键的那个“空”就是GPU上的计算程序。对于缺少CUDA编程经验的用户,可以尽量复用原来代码,避免新写CUDA kernel。


这里我们演示了如何把EDVR里面的Deformable Convolution包装成TensorRT plugin(代码在这里:
https://github.com/shining365/EDVR-TRT/blob/master/trt_onnx/DeformConvPlugin.h  )
。对于这个PyTorch的例子来说,我们尽量保持原始代码不变,原封不动地把相关代码片段提取出来,并拷贝了原始代码的编译选项,使得CUDA代码可顺利编译。


11 使用fp16/int8加速计算


如果模型已经成功地跑在了TensorRT上,可以考虑使用fp16/int8做进一步加速计算。TensorRT默认运行精度是fp32;TensorRT在Volta、Turing以及Ampere GPU上支持fp16/int8的加速计算。


使用fp16非常简单,在构造engine时设置标志即可。这一点体现在trtexec上就是它有--fp16选项,加上它就设置了这个标志。


我们举例说明fp16加速计算的重要意义。对于EDVR,用ONNX导出的模型,直接运行fp32加速比是0.9,比原始模型慢,但是打开fp16就有了1.8倍加速。fp16对精度的影响不是很大。


int8量化需要校正数据集,而且这种训练后量化一般会损失精度。如果对此介意,可以考虑使用Quantization Aware Training,做训练时量化。


12 发挥TensorRT的极致性能


前面讲的是TensorRT的一般用法,当你成为TensorRT熟手之后需要考虑如何发挥TensorRT极致性能。


1)API搭建网络

对于EDVR来说,我在TensorRT上用过两种方式运行,一种是用ONNX导出,它的fp32和fp16精度下的加速比是0.9和1.8;另一种是API搭建,它的加速比是1.1和2.7。可以看出API搭建有一定收益。假如模型特别重要,就要考虑用API搭建。


2)优化热点

通过Nsight Systems可以找到时间占用最多的操作,对它进行重点优化。


3)用Plugin手工融合所有可以融合的层

以上这些方面都做到的话,基本上就可以做到在TensorRT上的极致性能。


13 总结与建议


今天我们推荐的开发方法是用ONNX Parser导入模型。这里需要熟悉Graph Surgeon用法,针对各种特殊情况处理。有可能需要自定义Plugin,包装现有CUDA代码。我们推荐使用混合精度,特别是fp16用法简单、效果不错;int8有更好计算性能,但一般会有精度下降。如果想要进阶,要试着使用API搭建网络,并且编写与优化CUDA kernel。


14 示例代码


以上就是我分享的全部内容,谢谢。


直播回放:

https://www.livevideostack.cn/video/gary-ji/ 


- The cover from creativeboom.com


还可输入800
全部评论
作者介绍

季光

英伟达

DevTech团队技术负责人

文章

粉丝

视频

相关文章
阅读排行
  • 2周
  • 4周
  • 16周
热门视频

WebRTC视频数据流程分析

许建林/《WebRTC Native开发实战》书籍作者