一、前言

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优秀本质

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的