深度学习模型编译框架TVM概述

在任意深度学习的应用场景落地一个模型/算法时,需要经历两个基本步骤:1. 根据数据生产一个模型的训练步骤;2. 将生产出的模型部署到目标设备上执行服务的推理步骤。训练步骤目前基本由Tensorflow、PyTorch、Keras、MXNet等主流框架主导,同样的,推理步骤目前也处在“百家争鸣”的状态。

TVM是什么?

TVM是一款开源的、端到端的深度学习模型编译框架,用于优化深度学习模型在CPU、GPU、ARM等任意目标环境下的推理运行速度,常见的应用场景包括:

  • 需要兼容所有主流模型作为输入,并针对任意类型的目标硬件生成优化部署模型的场景

  • 对部署模型的推理延迟、吞吐量等性能指标有严格要求的场景

  • 需要自定义模型算子、自研目标硬件、自定义模型优化流程的场景

深度学习模型编译框架TVM概述_第1张图片

TVM软件框架(图片来自TVM官方网站)

TVM框架如上图:主流的深度学习框架(Tensorflow, Pytorch,MXNet等)导出的模型作为TVM框架的输入,经过该框架内一系列的图优化操作以及算子级的自动优化操作后最终转化为针对目标运行时(CPU/GPU/ARM等)的部署模型,优化后的模型理论上可以最大化地利用目标硬件的资源以最小化模型的推理延迟。

为什么用TVM优化模型推理?

模型推理场景下用于模型优化、部署的软件框架仍处于“百家争鸣”的状态,其原因在于推理任务的复杂性:训练后的模型需要部署于多样的设备上(Intel CPU/ NVGPU/ ARM CPU/FPGA/ AI芯片等),要在这些种类、型号不同的设备上都能保证模型推理的高效是一项极有挑战的工作。

一般来说,主流的硬件厂商会针对自家硬件推出对应的推理加速框架以最大化利用硬件性能,如Intel的OpenVINO、ARM的ARM NN、Nvidia的TensorRT等,但这些框架在实际应用场景中会遇到不少问题:

  • 厂商推理框架对主流训练框架产生的模型的算子种类支持不全,导致部分模型无法部署

  • 模型部署侧的开发人员需要针对不同的硬件编写不同的框架代码,花精力关注不同框架对算子的支持差异和性能差异等

因此,一套可以让我们在任意硬件上高效运行任意模型的统一框架就显得尤其有价值,而TVM正是这样一套框架。

TVM如何优化模型推理?

实 际 上,“运 行 模 型/代码 到 任 意 种 类 的 硬件”并不是一个概念上全新的课题。在计算机编程语言发展的早期阶段(第二代编程语言),人们也曾遇到过类似的困境,即一种硬件平台必须配套一种汇编语言且不同汇编语言无法跨平台运行的情况。随着该领域的发展,人们给出了解决之道——引入高级语言和编译器,如下图:

深度学习模型编译框架TVM概述_第2张图片汇编语言向高级语言-编译器路线发展

  • 程序员负责用高级语言描述上层业务逻辑,不必关注具体硬件特性

  • 编译器将高级语言逐层转化为更底层的符号,也称中间表示(IR),其中最底层的IR可以对接不同的硬件,进而转化为针对不同目标的机器码

TVM框架正是借鉴了这种思想,我们可以把TVM理解成一种广义的“编译器”:TensorFlow、PyTorch等训练框架导出的模型可以认为是“高级语言”,而TVM内部的图级别表达式树、算子级的调度Stages则可以认为是“高级语言”的“中间表示”,如下图。

深度学习模型编译框架TVM概述_第3张图片深度学习模型“编译器”

TVM框架构成

以“模型部署”为边界,TVM可以分为TVM编译器和TVM运行时两个组件:

深度学习模型编译框架TVM概述_第4张图片TVM组件:编译器和运行时(图片来自TVM官方网站)

编译器负责模型的编译和优化,是TVM的主体功能:

  • 编译优化过程支持Python和C++接口

  • 系统环境支持Linux、Windows以及MacOS平台(部分功能如AutoTVM在非Linux平台可能受限)

运行时负责在目标设备上执行编译器生成的模型推理代码:

  • 部署过程支持JS, Java, Python, C++语言

  • 部署平台除了支持Linux、Windows以及MacOS系统,还支持Android, IOS, 树莓派等端侧系统

TVM编译过程

TVM编译器是TVM的主体功能组件,负责优化、编译深度学习模型为可在目标设备运行推理任务的代码,整体编译过程如下图所示:

深度学习模型编译框架TVM概述_第5张图片TVM编译栈(图片来自TVM官方网站)

上图中,蓝色方框代表TVM编译过程中涉及到的数据结构,黄色方框代表处理这些数据结构的算法,粉色的AutoTVM则是辅助Schedule功能选择参数的一种特殊算法。

整个TVM编译栈的数据结构包括:

  • Model from Frameworks: 从PyTorch、MXNet等框架导出的模型

  • IR Module(Relay): TVM的图级别IR,数据结构为Relay Expression构成的Relay AST,查看源码可以发现,relay.Function实际上是relay.Expr的一种子类,其功能是为Relay Pass(遍历Relay AST的手段)提供一个入口点,故而官网中也把Relay IR称为"End to end function"

  • IR Module(TIR): TVM的Tensor级别IR,仍然以AST的形式组织Expr,但包含了各算子(对应网络层)的具体调度细节(循环嵌套、并行、指令集等),这一层IR由一系列的tir.PrimFunc组成,此处的PrimFunc不再是访问整个AST的“入口点”,而是对应神经网络的一个层或融合层(如Conv+BN+Activation)的具体计算调度过程

  • Runtime Module: TVM编译栈的最底层IR,直接对接Runtime在目标设备上运行

TVM编译栈涉及的数据处理算法包括:

  • Frontend:  负责将外部框架生成的神经网络模型转化为TVM的图级别模型Relay IR

  • Relay Passes: 负责在图级别优化IR的计算过程,常见的图优化Pass包括常量折叠、死代码消除、推理简化、内存布局转化、Op融合等

  • Scheduling: 负责转化Relay Passes优化后的Relay IR为TensorIR,大致的流程是对于每个层/融合层对应的算子,执行graph_runtime_codegen.cc的Codegen,根据TVM框架中注册的算子compute和schedule函数,将每个算子转化为具体的计算调度过程(Relay IR --> FTVMCompute --> Tensors --> FTVMSchedule --> Schedule Stages

  • TIR Passes: 负责Tensor级别的IR优化,常见的如access index simplification,也有负责IR lower的function entry decoratation等

  • Target Translation(CodeGen):  设备相关的低级别Codegen,将TIR转化为TVM Runtime所需的Module

TVM特色

特色1:算子的计算和调度过程定义的分离

TVM的一个特色是借鉴Halide语言的计算和调度分离定义的设计方法,将深度学习模型的算子计算方法定义和具体计算调度方式做了解耦:

•计算(Compute)定义了每个Relay Op的具体计算过程

•调度(Schedule)则细化了计算过程每个具体步骤的实际计算细节

以卷积算子为例,计算(Compute)定义的是算子应如何根据卷积核、卷积步长等参数操作输入Tensor,并输出卷积后的Tensor的过程;调度(Schedule)则定义了更具体的计算细节,例如卷积的6层计算循环的顺序、每层循环是否需要并行计算、并行数设定为多少等。

之所以做这样的分离式设计,是因为一般情况下,特定算子的计算方式是不变的,而算子的调度却有几乎无限种配置方法,显然对于不同的Target如CPU、GPU,其最优调度配置必然大相径庭;即使对于同一种Target如LLVM,不同型号的CPU也会导致最优配置相差甚远。TVM的计算-调度分离的设计方式,给二次开发人员提供了足够便利的API,方便其自定义性能更佳的调度配置。

特色2:AutoTVM

默认情况下,TVM为定义的所有算子针对常见的Target(x86/ARM/CUDA等)各提供了一套默认的调度配置。但一方面指定Target的默认配置不可能在所有型号的硬件下最大程度地利用运行环境的计算资源;另一方面调度的可选配置空间巨大,人工手动调整调度效率低且无法保证调度性能。为了解决该问题,TVM开发了AutoTVM模块通过搜索的方式获得具体硬件下的最优调度配置,总体步骤如:

  • 针对每个Relay子函数,定义一系列调度原语(Schedule primitives),通过调度原语的组合实现计算结果等价的Relay子函数

  • 在海量的配置组合中为每个Relay子函数搜索最优的调度配置

对于卷积这样的计算密集算子,调度配置的搜索空间通常有几十亿种选择,比较棘手的问题就是如何高效的在搜索空间中获得最优/局部最优解。常见的做法有两种:

  • 遍历所有可能的参数编译模型并评估

  • 评估部分参数模型,训练一个性能评估函数,用于指导搜索过程寻找更好的调度配置

由于搜索空间过于巨大,第一种方法需要花费的时间难以估量;AutoTVM采用的是第二种,用一个XGBoost模型(其他回归预测模型也可)作为评估调度配置性能的模型(CostModel),用以指导模拟退火算法(默认)寻找最优调度配置。

搜索过程完毕后AutoTVM会将过程中记录的最优配置保存下来,编译模型时可用于代替默认的Target调度配置。

来Adlik Github看看吧~

深度学习模型编译框架TVM概述_第6张图片

你可能感兴趣的:(算法,数据结构,大数据,编程语言,python)