深度学习模型部署TensorRT为何如此优秀?
一、前言
PyTorch模型的高性能部署问题,主要关注两个方面:高度优化的算子和高效运行计算图的架构和runtime。python有快速开发以及验证的优点,但是相比C++来说速度较慢而且比较费内存,一般高性能场景都是使用C++去部署,尽量避免使用python环境。
TensorRT为什么那么快,因为engine在构建的时候,在每个平台(A10、A100、T4等)上搜索到了最优最快的kernel(实现了一些op)。高效率运行计算图也是很关键的一点,TensorRT构建好engine后,需要libnvinfer.so来驱动,其中实现过程很容易猜到:
- 序列化和反序列化,也就是所谓的生成engine,读取engine;
- 推理engine、多stream运行计算图,管理engine所需要的一些环境,比如显存和中间变量等;
为了达到极致的性能,TensorRT的整个运行时都是在C++环境中,虽然提供了Python-API,但实际调用执行的操作都是在C++中,Python只提供包了一层的作用,算子和执行整个计算图的地方都是C++。 此外,Nvidia最清楚自家GPU或DLA该如何优化,所以TensorRT跑网络的速度是最快的,比直接用Pytorch快N倍。
二、深度分析 - TensorRT优秀本质
1. 搜索整个优化空间
与Pytorch等其它训练框架最大区别是,TensorRT的网络优化算法是基于目标GPU所做的推理性能优化,而其它框架一方面需要综合考虑训练和推理,更重要的是它们没有在目标GPU上做针对性的优化。那么,TensorRT是如何针对目标GPU优化的呢?
简单讲就是在可能的设计空间中搜索出全局最优解。
这个搜索空间有哪些变量呢?
比如CUDA架构中的编程模型所对应的,将tensor划分为多少个block?以及这些block如何组织到Grid中:
任务被划分为多个Block
Block以Grid的方式组织起来
不同的组织层次以对应不同的存储体系结构中的不同存储器
再举例,使用什么样的指令完成计算,可能是FFMA、FMMA,可能是TensorCore指令...
更难的部分可能是tensor数据流的调度,把他们放在local、share还是global memory呢?如何摆放呢?
这些变量组合在一起是一个巨大的搜索空间,可能你的CPU计算几天也得不出个结果来。
但是我们知道神经网络的计算是由一个个粒度更大的算子组成的,算子上面还有粒度更大的层结构。我们也清楚地知道层与层之间相对独立,也就是说可以针对每层计算优化,最后把优化后的层串在一起大概率就是网络的全局最优解。
于是,TensorRT预先写了很多算子和层(CUDA Kernel)。当然这些算子的输入和输出tensor是可以配置的,以适应网络输入和输出的不同以及GPU资源的不同。
部分优化好的算子
搜索空间变小了,从原来的指令级别的搜索,上升到了算子级别的搜索。因为这些实现都是用CUDA kernel所写,更准确的说是Kernel级别的搜索了。
但是tensor数据流的调度问题并没有解决,这也是最关键和复杂的地方。我们应该将输入tensor划分为多少个Block呢?这些Blocks应该分配给多少个线程呢?tensor存储在哪呢?local/share/global memory的哪些地方呢?中间计算结果存储在哪里呢?
对于计算部分是可以通过模拟的方式(类似指令集仿真器)计算得到性能的,但是tensor数据流在share/L2/Global Memory的流动过程就很难通过仿真计算得到精确结果,因为要被模拟的数据量和线程数过大,何况要尝试的可能性还很多,靠CPU仿真计算的思路就别想了。唯一办法就是让候选算子在目标GPU上直接跑跑,统计出性能,最后通过比对选出最优解。TensorRT把这个过程叫做Timing,TensorRT甚至可以将优化的中间过程存储下来供你分析,叫做timing caching(通过trtexec --timingCacheFile=<file>)。
Nvida GPU memory架构
以上所描述的优化过程可以叫做Hardware Aware Optimazation
总结起来优化器会重点分析:
- Type of hardware(Hardware capability...)
- Memory footprint(Share, Cache, Global...)
- Input and output shape
- Weight shapes
- Weight sparsity
- Level of quantization (so, reconsider memory)
而这些是Pytorch等框架不会去深入挖掘的,尤其是对存储系统的优化。
2. 强制选择Kernel
由于Block之间线程的运行顺序是随机的,CPU可能在向GDDR/HBM读写数据,甚至GPU的时钟频率也在随负载的变化而变化,这导致了不同系统运行环境下GPU的性能表现会有差异。这种差异也可能导致TensorRT Timing的最优解不是实际推理时的最优解,可能选择了次优的Kernel。
TensorRT提供了一个补救方法,就是强制指定选择某个Kernel实现,如果你很确信它是最优解的话。
TensorRT提供的API叫做AlgorithmSelector
3. Plugin
当然,你对自己设计的算子更有把握,可以自己写Kernel,然后指定使用它
不过更多情况下,是因为发现TensorRT不支持某个算子,你才被迫去写Kernel,毕竟CUDA编程不简单,何况性能还需要足够好。
4. cuBLAS和cuDNN
TensorRT安装指导要求你先安装CUDA SDK和cuDNN
CUDA SDK需要安装是显而易见的,因为TensorRT所调用的Kernel需要NVCC编译器来编译成Nvidia GPU的汇编指令序列啊
但是CUDA SDK中还有一个cuBLAS库也是被TensorRT所依赖的,我们知道C++库BLAS(Basic Linear Algebra Subprograms),它是针对CPU进行的线性代数计算优化,那么cuBLAS就是针对CUDA GPU开发的线性代数计算库,它的底层当然也就是用CUDA Kernel写成的。典型的矩阵乘法算子就可以直接调用cuBLAS了。
cuBLAS开发的很早,应该是CUDA生态最早的一批库了吧,但是随着深度学习的普及,Nvidia又在生态中加入了cuDNN库,它的层次更高,封装了到了网络层,所以其实TensorRT也可以直接调用优化好的cuDNN库中的Kernel?是也不是
TensorRT可以选择所谓Tactic(策略)来决定是使用TensorRT写的Kernel还是cuBLAS和cuDNN的
5. Tactic
TensorRT的Tactic能决定很多优化选项
例如,每次timing某个算子时需要平均的运行次数。缺省TensorRT会运行四次,以降低不确定性带来的误差,但这个次数是可以修改的。
还可以决定上面提到的Kernel库的选择,Plugin的选择,GPU时钟频率锁定等。
6. 量化
TensorRT当然具备网络量化能力,提供了将全网都量化到int8的隐性量化方式,也提供了插入Q/DQ Layer的显性量化方式。
混合量化是Nvidia做的很优秀的地方,这对于高效利用计算资源起到了重要作用,不过,此处不做过多描述。
7. 多应用推理和多卡推理
多卡或多节点才是Nvidia的杀手锏。