TensorRT是如何做到架构更快的？ _架构

本文经自动驾驶之心公众号授权转载，转载请联系出处。
一. TensorRT是什么？2016年Nvidia为自家GPU加速推理而提供的SDK，人们有时也把它叫做推理框架。
二. 为什么？只有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实现，如果你很确信它是最优解的话。