TVM: Deep Learning模型的优化编译器(强烈推荐, 附踩坑记录)

本文作者是阿莱克西斯，原载于知乎，雷锋网获得授权转载。

（前排提醒，本文的人文内容部分稍稍带有艺术加工，请保持一定的幽默感进行阅读）

关注我最近想法的同学应该知道我最近都在把玩 TVM，今天终于使用 TVM 得到了非常满意的结果，而专栏也很长时间没更新了，于是来安利 (水) 一篇。

本来可能用不到 TVM，项目其实进展的很顺利，我们初始的 tensorflow 模型在 android 端得到了满意的 latency，我也可以照常一边修炼我的仙, 继续和有奶大定律, 自由单子, Kan-Extension 等邪魔外道搏斗... 一边稳稳的推进项目进度。

无奈 scientist 一意孤行要上 Pytorch, 于是我们换了一个 Pytorch 模型...

先不说同样的 SSD 魔改模型，Pytorch 在 android 端比 tensorflow 整整慢了 5 倍，光是把 Pytorch 模型移植到 Android 上都让开发团队整整褪层皮 (Pytorch 对 Android 的支持简直为 0，tensorflow 的工程支持相对 pytorch 简直无敌)。而这时候已经花了好多时间，项目眼看要 delay....

头都炸了的我在打算手撸 OpenCL 调优之前，去问了下我们组的 CV 大神该怎么办，大神微微一笑，转身随风而去，只听云端传来 3 个字：「T~V~M~~~~~"

于是我就开始 TVM 的研究 (踩坑) 之路, 到今天为止终于把所有的路都踩平了之后，成功把我们的 Pytorch 模型用 Auto-TVM 调优成功且部署在了我们的 android 系统上，性能整整提高了 8 倍，比我们之前的 tensorflow 模型还要快。更重要的是，通过 TVM，我们的调优完全不 couple 与硬件和模型 Framework，就算以后换模型，换终端，或者哪天 scientist 想不开要换回 tensorflow 或是使用 MXNet 都无所谓，用 auto-TVM 自动调调就行了（只可惜了我的 Cuda C 编程调优都白学了）。

简单介绍下 Auto-TVM 的调优终端设备的用法

1. 你可以有很多手机平板设备，安装好 TVM RPC 这个 App 之后，可以在 App 里输入 Tracker 的 IP 和端口，进行设备注册 (另外输入一个设备 ID 来让 Auto-TVM tuning 程序找到)。

2. Tracker 是一个 Python 的程序，git clone TVM 之后，按教程编译好，就可以按这个教程启动 Tracker。

3. Auto-TVM tuning 程序也是一个 python 程序，它会连接 Tracker(也可以和 Tracker 是一台机器) 找到相应的设备 ID 的 IP，然后和设备直接用 RPC 通信，Auto-TVM 程序会根据程序预设的 target(比如是不是 arm cpu，要不要用 OpenCL...) 来把你想要优化的 Deep Learning 模型直接编译为设备的 machine code, 通过 TVM RPC 把 code 部署在终端，终端的 TVM RPC App 会测试这个模型的 inference performance，然后回报给 Auto-TVM tuning 程序，然后 Auto-TVM tuning 程序会根据反馈，重新计算该如何优化编译，重新生成新的模型的 machine code 再次部署... 如此循环... 直到达到预设的实验次数 (比如 2000), 或太多次实验都没有提高提前结束 (比如第一次就找到了最优优化结果)。最后 TVM 会根据调优时得到的最佳「编译参数」来最终编译你的 deep learning 模型为终端模型的 machine code，最终完成优化编译过程。

以上只是简单介绍，具体请看 TVM 的论文，和去 TVM 官网看 tutorial，写得非常详细切提供了很多很好理解的范例代码。我的最终的 tuning 程序，就是魔改其中一个范例程序而来。

TVM 踩坑记录

TVM 目前还只是 0.6 版本，很多东西还不稳定，由于开发环境各异，有时候需要工程师自己解决一些开发团队在开发时没有碰到的问题，但是这些问题相对与 TVM 提供的巨大优势相比，都是小问题啦（工程能力越强，魔改力越强，你就可以越早体验新技术带来的好处呀。）。（我遇到的最坑的问题其实是公司网络各种 IP 禁止访问，封端口，使得 android 机和开发服务器一直连不上, 最终还是在自己的电脑上装了虚拟机，自建了一个小 LAN 才解决的这个问题）

1.编译 tvm4j-core 出错: cannot find symbol [ERROR] symbol: class SharedSecrets

JDK11 会遇到这个问题，因为 JDK11 已经把 sun.misc.SharedSecrets 换到别的地方了，建议不要尝试修改 TVM 源代码来 fix 这个问题，因为你会遇到其他更多问题，请下载 JDK8，把 JAVA_HOME 设为 JDK8 的，一切就会很顺利

2.Android TVM RPC 编译出错: #error "Unable to determine endianness of your machine; use CMake to compile"

Android RPC server fails to build

按上边 link 里的修改 endian.h 文件即可，参见我下边的修改

diff --git a/include/dmlc/endian.h b/include/dmlc/endian.h

index 5bf53fb..9422fce 100644

--- a/include/dmlc/endian.h

+++ b/include/dmlc/endian.h

@@ -23,7 +23,9 @@

#elif defined(__EMSCRIPTEN__)

#define DMLC_LITTLE_ENDIAN 1

#else

- #error "Unable to determine endianness of your machine; use CMake to compile"

+ #include

+ #define DMLC_LITTLE_ENDIAN (__BYTE_ORDER == __LITTLE_ENDIAN)

+ /*!#error "Unable to determine endianness of your machine; use CMake to compile" */

#endif

#endif

3.Auto-TVM 运行时出错"Do not know how to handle return type code 113"

Auto-TVM failed on Android Device, with error msg of "Do not know how to handle return type code 113"

可以根据我上边在 TVM Discussion 里的自问自答来解决。

4.找不到 TVM_NDK_CC

[SOLVED] Android_rpc_test.py failed

按照 dayanandasiet 的回复设定 TVM_NDK_CC 即可

Follow the below steps to generate toolchian and try to generate application with below export

comand

Tool chain generate with below instruction

./make-standalone-toolchain.sh --platform=android-24 --use-llvm --arch=arm64

--install-dir=/home/user/software/android-toolchain-arm64/

Download Java and SDK, set proper path

export TVM_NDK_CC=/home/user/software/android-toolchain-arm64/bin/aarch64-

linux-android-g++

export JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64

export ANDROID_HOME=/home/user/software/android-sdk-linux

build mxnet model with nnvm with below config/parameter and use same library,

param and graph on your android application

target =‘llvm -target=arm64-linux-android’

target_host = None

reference mobile_darknet_save.py 2

Compile application android deploy 1 using this config.mk 2configuration for CPU flavor

5.LLVM only Large Small are allowd on AArch64

https://github.com/dmlc/tvm/issues/2005 可解。

6.Auto-TVM 自动优化时出错：Cannot find config for target=cuda

这个不是什么大问题，某 operator 不能调，对我来说其他的可以调就行了。。。。

7.Auto-TVM 自动优化 OpenCL 时出错: No OpenCL platform matched given existing options

No OpenCL platform matched given existing options

也是自己问最终自己解决的，很反直觉，编译 TVM 的时候，选择 OpenCL=OFF，就没有这个问题，选择 OpenCL=ON，为终端 Cross Compile OpenCL 就不 work 了... 应该是 bug.

8.Auto-TVM 自动优化 OpenCL 时出错: CL_INVALID_WORK_GROUP_SIZE

CL_INVALID_WORK_GROUP_SIZE error after auto-tuning for OpenCL on Android Device

应该是我 trial number 设的太小了，以至于 TVM 找不到一个 valid 的 kernel，顺着这个问题，发现了 CL_INVALID_WORK_GROUP_SIZE 的一个 undocumented 的错误源，即 OpenCL kernel 使用过多的 register file 也会造成 CL_INVALID_WORK_GROUP_SIZE 错误，这一点在查 OpenCL 文档的时候是查不到的, 有点 tricky。

9.Auto-TVM 自动优化时 Android TVM RPC Crush，一切白调。。。

目前 TVM 还不支持 checkpoint，不过我们可以很简单的魔改 measure_methods.py 这个文件，把 190 行 set_task: 这个函数里的 check remote 失败直接 raise RuntimeError 退出程序，改成循环多次 check 即可，这样使得 Auto-TVM 一方持续等待 Android 程序上线，比一点网络问题，或者终端问题，就废掉之前 n 多个小时的 auto-tuning 成果要好的多。

最后感谢

@ 陈天奇

大神为我们带来了这么方便的工具。

好了，今天就到这里，我继续修仙炼丹去了～