TVM工具栈

这两天要回归TVM工具栈的社群了,在这里写点东西介绍一下。

现在做数据分析的都叫数据科学家吧。虽然名头是个“科学家”,到头来还是个工程师。工程和理论的不同之处就在于,工程是面向现实的。把理论中的各种假设摘掉,再把经济问题混进去,不管什么问题都会变得异常棘手。神经网络也逃不出工程这个怪圈,一系列NN特色的问题出来了:

  • MxNet、PyTorch、Caffe、TensorFlow……这么多框架,要是每个框架训练出来的模型都要放上线,岂不是每个框架都要装?
  • 我MxNet训练出来的网络怎么放Windows Server上?(比较早的时候MxNet不支持Windows)
  • 我要在手机的ARM片子上运行神经网络,训练用的框架没有对这种架构进行优化,跑起来很慢

谷歌说,我要用自己的方式解决问题,于是有了给移动端用的TensorFlow Lite和给服务器用的TensorFlow Serving。甚至还弄出了TensorFlow.js这种东西给浏览器跑网络……不过谷歌的TF怎么说还是一个闭环,要用TF的工具,就得一整套全用了:训练用TF,部署用前述的那些,至于别家框架训练出来的东西,还得费点力迁移一下。

其他家的框架看TF吃独食不顺眼,纷纷站队要共荣。于是Facebook和微软义结金兰,合作出了一个通用的模型格式ONNX,于是PyTorch、Caffe2、CNTK训练出来的模型可以通用了。

虽然ONNX统一了模型参数的储存格式,但以前训练的模型还得要啊。于是框架普遍都选择另造工具来把模型文件转换成ONNX,而非直接把模型格式换成ONNX——包括PyTorch和Caffe2,这两个框架仅仅是在软件里内建了ONNX的存取程序,所以目前唯一直接支持ONNX只有软家刚出生的CNTK。不管怎么样,有了标准格式总算是能共享模型数据了。

当然了,做神经网络要操心的不仅仅是用的什么框架,也要头疼该在什么东西上面执行计算图。像刚才说的ARM平台,在还没有专门的神经网络执行单元(NPU)的时候,执行效率最高的是腾讯的NCNN(虽然是员工的个人作品),是靠手写Neon汇编做的优化。虽然NCNN也能在x86架构的处理器上使用,但是效率显然达不到我们所能在ARM设备上期待的水平。想想看,除了平常的CPU配上SIMD优化,神经网络的显卡接口就有NVIDIA的CUDA(cuDNN)、AMD的ROCm、OpenCL和OpenGL、Vulkan两个非主流,更别提FPGA那种幺蛾子……框架的维护者不但要花很多时间适配新冒出来的平台,对不同平台又要使用不同的优化策略,工作量一下就大了好多——这也就算了,反正烦维护的人不烦我……但是要是框架没有提供接口给自己做业务用的语言怎么办?就算有接口,框架那么大一个,直接塞用户手机里是不是有点过分?但是如果用最小生成体积的方法,手写网络运算图的话,自己的工作量又太大了。要是能直接用既有模型直接生成这个网络的运算程序就好了。

看到这里相信你已经差不多知道TVM这一套杂技是解决什么问题的了:NNVM把各个前端——也就是框架——训练出来的模型转换成TVM使用的通用格式,然后TVM将模型文件编译成不依赖任何框架、在指定后端设备(如CUDA)上执行的二进制库。同时因为TVM掌管了库代码的生成,TVM可以对代码进行优化。

举例来说,如果在CUDA核上运行,框架每执行到一个网络层就调用一次GPU。因为框架的计算程序都是写好的,像加法这种操作,在框架的代码里写死了就只有两个参数。三个层相加这种操作,框架必须调用两次GPU,中间还要产生一个不必要的内存块储存中间量。浪费内存不说,CPU和GPU同步就浪费了不少时间。框架的执行流程就像这样:

A' = A1 + A2 // 第一次CUDA调用,对两个张量求和,保存中间值。
B = A' + A3 // 第二次CUDA调用,对两个张量求和。

而TVM自己就管代码生成,能够直接优化这两步运算为一步:

B = A1 + A2 + A3 // 唯一一次CUDA调用,直接对三个张量求和。

而且因为TVM不会生成网络里没有用到的东西,没有用到乘法就不会生成任何乘法相关的代码,最后编译出来的运算库非常小,非常适合在各种设备上进行部署——毕竟TVM本来就是为了部署而生的嘛。

顺便一提,虽然没有仔细看过……TVM好像真的可以生成FPGA用的Verilog代码……