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

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上还带有硬件编码器，可以将处理后的图片进行编码，输出视频码流。在上述流程中，无论是解码，还是数据的处理，还是最后的编码，都可以使