tvm部署c++神经网络前向代码到android端

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 
#include 
#include 
#include 

#include 
#include 
#include 

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(json_in)), std::istreambuf_iterator());
    json_in.close();

    // parameters in binary
    std::ifstream params_in("deploy.params", std::ios::binary);
    std::string params_data((std::istreambuf_iterator(params_in)), std::istreambuf_iterator());
    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(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(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目录下。