tvm部署c++神經網絡前向代碼到android端
tvm是端到端的神經網絡編譯器。簡單來說，他可以把神經網絡模型編譯成一個動態鏈接庫，並部署到各種硬件上去執行，包括移動設備。
模型
tvm可以支持編譯各種框架模型，包括tflite，onnx等，本文主要描寫從onnx到tvm部署到android的過程。
tvm支持的框架以及各種模型編譯教程網址：https://docs.tvm.ai/tutorials/index.html
生成交叉編譯工具
我實在ubantu上編譯onnx，需要下載android sdk，並利用sdk生成一個交叉編譯器。因爲我是在ubantu平臺上編譯arm平臺的鏈接庫。
cd /sdk/ndk-bundle/build/tools/
./make-standalone-toolchain.sh --platform=android-24 --use-llvm --arch=arm64 --install-dir=/opt/android-toolchain-arm64
我用的架構師arm64-v8a，所以–arch參數爲arm64。最後生成的交叉編譯器在/opt/android-toolchain-arm64目錄下。
tvm下載和編譯
1.克隆倉庫
git clone --recursive https://github.com/dmlc/tvm
2.安裝依賴
sudo apt-get update
sudo apt-get install -y python python-dev python-setuptools gcc \
libtinfo-dev zlib1g-dev build-essential cmake
3.安裝llvm
要安裝大於4.0版本的，而ubuntu 16.04 apt官方源最新只有3.x，ubuntu 18.04則沒問題（安裝的是6.0）。如果apt官方最新的llvm版本小於4,那麼使用llvm的源：
apt install software-properties-common
apt-add-repository "deb http://apt.llvm.org/xenial/ llvm-toolchain-xenial main"
apt-get update
apt-get install clang
這裏最後一步如果報了***文件找不到的錯誤，可以touch *** 新建一個空文件試一下，可能會有四到五文件報錯說找不到，新建一個空文件就可以。
4.移動config.cmake
cd tvm
mkdir build
cp cmake/config.cmake build
編輯build/config.cmake文件，裏面有一些功能開關，打開了set(USE_LLVM ON).
5.編譯
（1）編譯生成x86-64版本的tvm
make clean
make -4j
由於是在ubantu環境下，默認編譯x86-64版本的tvm。
（2）編譯生成arm64-v8a版本的tvm
首先導入環境變量
export AR_host="ar"
export CC_host="gcc"
export CXX_host="g++"
export LINK_host="g++"
export ARCH=arm64
export PATH=/opt/android-toolchain-arm64/bin:$PATH
export CROSS_COMPILE=aarch64-linux-android-
export CC=/opt/android-toolchain-arm64/bin/aarch64-linux-android-gcc
export CXX=/opt/android-toolchain-arm64/bin/aarch64-linux-android-g++
export LD=/opt/android-toolchain-arm64/bin/aarch64-linux-android-ld
export AR=/opt/android-toolchain-arm64/bin/aarch64-linux-android-ar
export AS=/opt/android-toolchain-arm64/bin/aarch64-linux-android-as
export RANLIB=/opt/android-toolchain-arm64/bin/aarch64-linux-android-ranlib
顯然，這裏引入的是之前的交叉編譯工具：/opt/android-toolchain-arm64
然後執行：
make clean
make -4j
問題：這裏爲什麼要編譯兩種版本的tvm？
在部署到移動端的過程中，兩個版本的tvm都要用到。在將onnx編譯成arm64版本的動態鏈接庫的時候，要用到x86-64版本的tvm，因爲我們是在x86-64的宿主機上進行的編譯；而arm64版本的tvm裏面有一些動態鏈接庫我們要放到移動端。
編譯過後會生成如下的動態鏈接庫文件，只是兩種編譯方式生成的so文件架構不同。
[ 5%] Linking CXX shared library libvta.so
[ 12%] Linking CXX shared library libtvm_runtime.so
[ 86%] Linking CXX shared library libtvm.so
[ 94%] Linking CXX shared library libtvm_topi.so
[100%] Linking CXX shared library libnnvm_compiler.so
編譯onnx
編譯onnx需要在x86-64版本的tvm下進行，需要安裝onnx：pip installl onnx
設置python環境變量：
export TVM_HOME=~/tvm/
export PYTHONPATH=$TVM_HOME/python:$TVM_HOME/topi/python:$TVM_HOME/nnvm/python:${PYTHONPATH}
編譯代碼：
import onnx
import numpy as np
import tvm
import tvm.relay as relay
import os
from tvm.contrib import util, ndk, graph_runtime as runtime
from tvm.contrib.download import download_testdata
onnx_model = onnx.load('****.onnx')
x = np.ones([1,3,256,256]) //輸入的tensor shape
arch = "arm64"
target = "llvm -target=%s-linux-android" % arch //編譯的目標架構
input_name = 'input' //網絡輸入節點名
shape_dict = {input_name: x.shape}
sym, params = relay.frontend.from_onnx(onnx_model, shape_dict)
with relay.build_config(opt_level=0):
intrp = relay.build_module.create_executor('graph', sym, tvm.cpu(0), target)
dtype = 'float32'
with relay.build_config(opt_level=0):
graph, lib, params = relay.build_module.build(sym, target, params=params)
print("Output model files")
libpath = "model.so"
lib.export_library(libpath, cc="/opt/android-toolchain-arm64/bin/aarch64-linux-android-gcc")
graph_json_path = "model.json"
with open(graph_json_path, 'w') as fo:
fo.write(graph)
param_path = "model.params"
with open(param_path, 'wb') as fo:
fo.write(relay.save_param_dict(params))
以上代碼生成三個文件model.so, model.json, model.params。這三個文件要放到安卓的assets目錄下。
移動端部署
其實移動端部署的C++代碼官網已經寫的非常清楚：https://docs.tvm.ai/deploy/nnvm.html
將以下C++文件命名爲libnative-lib.cpp
#include <dlpack/dlpack.h>
#include <tvm/runtime/module.h>
#include <tvm/runtime/registry.h>
#include <tvm/runtime/packed_func.h>
#include <fstream>
#include <iterator>
#include <algorithm>
int main()
{
// tvm module for compiled functions
tvm::runtime::Module mod_syslib = tvm::runtime::Module::LoadFromFile("deploy.so");
// json graph
std::ifstream json_in("deploy.json", std::ios::in);
std::string json_data((std::istreambuf_iterator<char>(json_in)), std::istreambuf_iterator<char>());
json_in.close();
// parameters in binary
std::ifstream params_in("deploy.params", std::ios::binary);
std::string params_data((std::istreambuf_iterator<char>(params_in)), std::istreambuf_iterator<char>());
params_in.close();
// parameters need to be TVMByteArray type to indicate the binary data
TVMByteArray params_arr;
params_arr.data = params_data.c_str();
params_arr.size = params_data.length();
int dtype_code = kDLFloat;
int dtype_bits = 32;
int dtype_lanes = 1;
int device_type = kDLCPU;
int device_id = 0;
// get global function module for graph runtime
tvm::runtime::Module mod = (*tvm::runtime::Registry::Get("tvm.graph_runtime.create"))(json_data, mod_syslib, device_type, device_id);
DLTensor* x;
int in_ndim = 4;
int64_t in_shape[4] = {1, 3, 256, 256};
TVMArrayAlloc(in_shape, in_ndim, dtype_code, dtype_bits, dtype_lanes, device_type, device_id, &x);
// load image data saved in binary
std::ifstream data_fin("cat.bin", std::ios::binary);
data_fin.read(static_cast<char*>(x->data), 3 * 256 * 256 * 4);
// get the function from the module(set input data)
tvm::runtime::PackedFunc set_input = mod.GetFunction("set_input");
set_input("data", x);
// get the function from the module(load patameters)
tvm::runtime::PackedFunc load_params = mod.GetFunction("load_params");
load_params(params_arr);
// get the function from the module(run it)
tvm::runtime::PackedFunc run = mod.GetFunction("run");
run();
DLTensor* y;
int out_ndim = 2;
int64_t out_shape[2] = {1, 1000};
TVMArrayAlloc(out_shape, out_ndim, dtype_code, dtype_bits, dtype_lanes, device_type, device_id, &y);
// get the function from the module(get output data)
tvm::runtime::PackedFunc get_output = mod.GetFunction("get_output");
get_output(0, y);
// get the maximum position in output vector
auto y_iter = static_cast<float*>(y->data);
auto max_iter = std::max_element(y_iter, y_iter + 1000);
auto max_index = std::distance(y_iter, max_iter);
std::cout << "The maximum position in output vector is: " << max_index << std::endl;
TVMArrayFree(x);
TVMArrayFree(y);
return 0;
}
這裏的頭文件是之前編譯tvm之後生成的，放到項目中src/main/cpp目錄下，和native-lib.cpp代碼一個目錄。
當然，這裏我項目裏的話其實是在jni層，這裏只說明tvm前向的語法，不會對jni的語法做描述，大家可以根據自己的需要將前向代碼移植到自己的jni層，然後可以在java層做調用。
注意C++端TVM的運行依於一個tvm的動態鏈接庫：arm64版本的libtvm_runtime.so，需要將該文件放到android項目的app/src/main/libs/arm64-v8a目錄下。