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 成果要好的多。

最後感謝

@ 陳天奇

大神為我們帶來了這麼方便的工具。

好了，今天就到這裡，我繼續修仙煉丹去了～

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！