优tech分享｜开源推理框架TNN模型部署加速与优化

人工智能的浪潮正席卷全球，与之而来的就是数据量快速上涨，运算力大幅提升，机器学习算法频频出现……而深度学习作为机器学习的重要分支，也开始逐步走到了聚光灯之下，迎来了蓬勃发展的局面，逐渐渗透进广泛的场景中，例如图像分类、目标检测分割、自然语言处理等都有它的身影。

在发展的同时，深度学习也存在着诸如模型计算量和OP类型都在急速增长，模型落地将面临性能、内存等一系列问题。对此业界涌现了ncnn/TNN等开源推理框架，尝试给出各自的解决方案。TNN作为腾讯优图实验室新一代的深度学习推理框架，支持移动端、服务端、桌面端，适配CPU、GPU、NPU等底层硬件，能实现人工智能在多个平台的轻量部署落地。

开源推理框架TNN究竟暗藏着什么秘密，它针对PyTorch模型部署的加速方案TNN-Torch又是怎样的，除此之外TNN还有哪些部署案例？近日，腾讯优图实验室高级研究员、高性能计算硕士、TNN项目成员dandiding以「开源推理框架TNN模型部署加速与优化」为主题，为大家进行了一次详细讲解。

深度学习已经发展很久，最近一次爆发从ALEXNET开始，大概每3至4月增长一倍，这对训练机器产生很大算力需求。2017年被大家称为AI芯片元年，Google TPU率先发布，号称领先GPU 10倍。随后，从APU到ZPU的上百款专用芯片陆续推出，GPU 也加入专用TensorCore指令，AI芯片战争全面打响, DSA遍地开花。

对于矩阵乘法来说，想要实现处理器的Peak计算性能，需要让处理器的计算单元都保持繁忙，例如每一个时钟周期都执行有效的FMA，每个Core的ALU持续不断计算。所以，在性能优化中，比较重要的事情是将数据“喂给”计算单元，使它能够持续不断地做计算任务。

第一步，对K维度的拆解，先将K控制在可控范围内；第二步，对M维度的拆解；第三步，对N维度的拆解。当完成三个维度拆解以后，可以得到较小尺寸，数据已经存留在Cache，这时候要进入更深入的Register Blocking

16×6和24×4是当前X86 AVX指令集上常用的分块大小。因为向量寄存器数量有限，只有16个，这样的分块可以让相对比较多的数据存留在寄存器。16×6计算密度更好，但24×4使用指令数更少。最终Microkernel的延迟，在不同体系架构上也可能会不同。

以线程尺寸为 4x4和8×8为例来说明不同尺寸的区别。4x4是load8个数据，做16次累加，8x8是load 16个数据，做64次累加。8x8的计算密度会更高一些，4x4对shared memory压力更大，但因为占用寄存器更少，可以实现更多的active warps，最终对global memory latency的掩盖效果更好。

现在业界有很多自动化方法，这里以Apache TVM和PyTorch TC为例来做介绍。TVM是一个端到端的DL compiler，它最核心的算法是把思想和调度做分离。TC最核心的区别是使用Polyhedral模型。

但是，TC目前没有办法实现性能最优，对计算密集算子优化的效果仍不理想。

TVM初版继承了Halide的思想，算法描述只说明做计算的逻辑，由调度语言描述具体如何完成这个计算过程。

TVM是端到端的编译, 他需要把前端TF、Pytorch等框架的模型，转到Relay IR，再经过图优化、调度搜索器等，最后由LLVM生成设备代码。搜索的过程加入了机器学习做耗时预测，加速搜索。TVM可以有效优化计算密集型算子，并可以通过融合策略策略打破算子边界。

TVM过往的 Layout策略：通过对算子按照影响程度分级，可划分为layout无关、layout弱相关、layout强相关三种， 由策略选定算子Format。

这样导致的结果是设计上对储存格式（NHWC、NCHW等）考虑不全面。于是改进点为：

01

扩充Conv，FC等算子的不同format实现；

02

Format为参数，加入搜索空间；

03

Tuning时，对不同格式做搜索，记录各Format时间；

04

以动态规划思路递推各格式耗时。

其次是对GPU调度优化。基于TNN 手工优化经验，优化扩充AutoTVM TOPI。主要包括：

01

扩充NCHWc Format支持；

02

优化 fp16 CUDA Core 实现；

03

优化TensorCore 调度，以Implicit GEMM的方式，在hw轴拆解为M；

04

通过padding，提高Scheduler参数覆盖。

目前存在模型转换链路长，Op缺一不可的现状，极易导致模型无法落地。

针对这种情况，需要减少模型转换失败率，让API尽量简单。比如在开发SDK时，做到预处理的尽量简单。对于静态维度Op，去提高它的性能。对于动态维度Op、自定义Op部分，能够满足短期上线的需求，待后续再做优化。

TNN-Torch的思路是参考了TRTorch做子图的思路，优化了转换子图的逻辑，并提供python和C++API。具体，腾讯优图做的优化主要是三个流程。

01

加入常量前馈操作，避免冗余节点。做子图分块时基于Op转换器的实现程度，以及它对不同Op的不同参数的实现程度做子图切分，保证切分后的子图在TN里面可以正常正确。

02

通过Op标记切割子图，并进行参数检验。可以让最终模型里面子图使用TNN加速，借助TNN多个平台的加速能力，模型可以在多个平台实现加速。

03

以ZeroCopy的方式做多个框架间Tensor传输的衔接。

对于一些应用模型，C++重写可以取得最佳的性能，但花费时间较多。通过TNN-Torch方案，模型很快可以上线，性能也可以达到不错的程度，更能提升大家的推理体验。

腾讯优图在Mac上面做CPU的人脸检测配准的DEMO，首先需要下载clone TNN的代码，提前准备下载依赖的模型，并将依赖的工具（比如opencv和protobuf）安装好。其次，cd 到对应的examples目录里面，打开webcam开关（其默认为关闭），然后再做编译，最后运行就可以，可以实现把人脸配准结果实时显示在图上面。

目前在TNN里面是已经开源，但没有正式上线，现在可以提供镜像测试。需要做的事情大致可分为环境准备、模型准备（使用torchvision里的maskrcnn模型 ）、数据准备（使用pytorch tutorials里面的图片）、编译、运行、结果查看、结果检验，整体的运行就结束了。相比直接Pytorch 会有80%+的性能提升。