在CPP下使用TVM来部署mxnet模型（以Insightface为例）

自从AI被炒作以来，各个深度学习框架层出不穷。我们通常来讲，作为AI从业者，我们通常经历着标注-训练-部署的过程。其中部署是较为痛苦的工作，尤其是在跨平台如（移动端需要native对接的时候。）当然用于inference框架同样也是层出不穷。但是大多数框架框架往往性能都一般，或者缺少相关op，或者就是转换模型较为困难。TVM的出现很大程度上为模型部署带来了福音。

但是网上将用于TVM部署的教程还比较少，尤其是通过cpp和移动端部署。本文以Insightface Model Zoo中的MobileFaceNet为例，介绍一下如何编译Mxnet模型、在python下inference、在cpp下inference、对比人脸余弦距离、以及在android下的部署。

安装

TVM编译环境的安装需要LLVM编译器，可以简要遵循官方的教程。 official installation tutorial.

LLVM 7.0 可能会导致编译错误，推荐使用LLVM 6.0.1

编译模型

TVM使用了一系列的优化措施来优化计算图，当模型编译完之后会生成若干个编译好的文件。在编译前要指定预编译的平台、架构、指令集等参数。

  1. import numpy as np
  2. import nnvm.compiler
  3. import nnvm.testing
  4. import tvm
  5. from tvm.contrib import graph_runtime
  6. import mxnet as mx
  7. from mxnet import ndarray as nd
  9. prefix,epoch = "emore1",0
  10. sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
  11. image_size = (112, 112)
  12. opt_level = 3
  14. shape_dict = {'data': (1, 3, *image_size)}
  15. target = tvm.target.create("llvm -mcpu=haswell")
  16. # "target" means your target platform you want to compile.
  18. #target = tvm.target.create("llvm -mcpu=broadwell")
  19. nnvm_sym, nnvm_params = nnvm.frontend.from_mxnet(sym, arg_params, aux_params)
  20. with nnvm.compiler.build_config(opt_level=opt_level):
  21.    graph, lib, params = nnvm.compiler.build(nnvm_sym, target, shape_dict, params=nnvm_params)
  22. lib.export_library("./deploy_lib.so")
  23. print('lib export succeefully')
  24. with open("./deploy_graph.json", "w") as fo:
  25.    fo.write(graph.json())
  26. with open("./deploy_param.params", "wb") as fo:
  27.    fo.write(nnvm.compiler.save_param_dict(params))

运行该代码后会生成三个文件分别为deploy_lib.so 、deploy_graph.json 、deploy_param.params 。其中deploy_lib.so 为编译好的动态库，deploy_graph.json为部署使用的计算图、deploy_param.params为模型参数。

使用TVM Python Runtime 进行简单的测试

TVM的Runtime(运行时)并不需要任何依赖，直接clone tvm后 make runtime.即可。

  1. import numpy as np
  2. import nnvm.compiler
  3. import nnvm.testing
  4. import tvm
  5. from tvm.contrib import graph_runtime
  6. import mxnet as mx
  7. from mxnet import ndarray as nd
  9. ctx = tvm.cpu()
  10. # load the module back.
  11. loaded_json = open("./deploy_graph.json").read()
  12. loaded_lib = tvm.module.load("./deploy_lib.so")
  13. loaded_params = bytearray(open("./deploy_param.params", "rb").read())
  15. input_data = tvm.nd.array(np.random.uniform(size=data_shape).astype("float32"))
  17. module = graph_runtime.create(loaded_json, loaded_lib, ctx)
  18. module.load_params(loaded_params)
  20. # Tiny benchmark test.
  21. import time
  22. for i in range(100):
  23.    t0 = time.time()
  24.    module.run(data=input_data)
  25.    print(time.time() - t0)

使用C++来推理MobileFaceNet人脸识别模型

在C++下 TVM Runtime（运行时）仅仅需要编译时输出的so文件，包含 “tvm_runtime_pack.cc” 。runtime的体积也比较小，只有几百K。

下列的CPP代码包含了通过输入一张对齐后的人脸识别照片，输出归一化的之后的人脸向量。

  1. #include <stdio.h>
  2. #include <opencv2/opencv.hpp>
  3. #include <tvm/runtime/module.h>
  4. #include <tvm/runtime/registry.h>
  5. #include <tvm/runtime/packed_func.h>
  7. class FR_MFN_Deploy{
  9.     private:
  10.         void * handle;
  12.     public:
  13.         FR_MFN_Deploy(std::string modelFolder)
  14.         {
  16.             tvm::runtime::Module mod_syslib = tvm::runtime::Module::LoadFromFile(modelFolder + "/deploy_lib.so");
  17.             //load graph
  18.             std::ifstream json_in(modelFolder + "/deploy_graph.json");
  19.             std::string json_data((std::istreambuf_iterator<char>(json_in)), std::istreambuf_iterator<char>());
  20.             json_in.close();
  22.             int device_type = kDLCPU;
  23.             int device_id = 0;
  24.             // get global function module for graph runtime
  25.             tvm::runtime::Module mod = (*tvm::runtime::Registry::Get("tvm.graph_runtime.create"))(json_data, mod_syslib, device_type, device_id);
  26.             this->handle = new tvm::runtime::Module(mod);
  28.             //load param
  29.             std::ifstream params_in(modelFolder + "/deploy_param.params", std::ios::binary);
  30.             std::string params_data((std::istreambuf_iterator<char>(params_in)), std::istreambuf_iterator<char>());
  31.             params_in.close();
  33.             TVMByteArray params_arr;
  34.             params_arr.data = params_data.c_str();
  35.             params_arr.size = params_data.length();
  36.             tvm::runtime::PackedFunc load_params = mod.GetFunction("load_params");
  37.             load_params(params_arr);
  38.         }
  40.         cv::Mat forward(cv::Mat inputImageAligned)
  41.         {
  42.             //mobilefacnet preprocess has been written in graph.
  43.             cv::Mat tensor = cv::dnn::blobFromImage(inputImageAligned,1.0,cv::Size(112,112),cv::Scalar(0,0,0),true);
  44.             //convert uint8 to float32 and convert to RGB via opencv dnn function
  45.             DLTensor* input;
  46.             constexpr int dtype_code = kDLFloat;
  47.             constexpr int dtype_bits = 32;
  48.             constexpr int dtype_lanes = 1;
  49.             constexpr int device_type = kDLCPU;
  50.             constexpr int device_id = 0;
  51.             constexpr int in_ndim = 4;
  52.             const int64_t in_shape[in_ndim] = {1, 3, 112, 112};
  53.             TVMArrayAlloc(in_shape, in_ndim, dtype_code, dtype_bits, dtype_lanes, device_type, device_id, &input);//
  54.             TVMArrayCopyFromBytes(input,tensor.data,112*3*112*4);
  55.             tvm::runtime::Module* mod = (tvm::runtime::Module*)handle;
  56.             tvm::runtime::PackedFunc set_input = mod->GetFunction("set_input");
  57.             set_input("data", input);
  58.             tvm::runtime::PackedFunc run = mod->GetFunction("run");
  59.             run();
  60.             tvm::runtime::PackedFunc get_output = mod->GetFunction("get_output");
  61.             tvm::runtime::NDArray res = get_output(0);
  62.             cv::Mat vector(128,1,CV_32F);
  63.             memcpy(vector.data,res->data,128*4);
  64.             cv::Mat _l2;
  65.             // normlize 
  66.             cv::multiply(vector,vector,_l2);
  67.             float l2 =  cv::sqrt(cv::sum(_l2).val[0]);
  68.             vector = vector / l2;
  69.             TVMArrayFree(input);
  70.             return vector;
  71.     }
  73. };

我们可以通过输入两张对齐后的人脸照片来提取人脸向量。

  1. cv::Mat A = cv::imread("/Users/yujinke/Desktop/align_id/aligned/20171231115821836_face.jpg");
  2. cv::Mat B = cv::imread("/Users/yujinke/Desktop/align_id/aligned/20171231115821836_idcard.jpg");
  3. FR_MFN_Deploy deploy("./models");
  4. cv::Mat v2 = deploy.forward(B);
  5. cv::Mat v1 = deploy.forward(A);

测量余弦相似度

  1. inline float CosineDistance(const cv::Mat &v1,const cv::Mat &v2){
  2.     return static_cast<float>(v1.dot(v2));
  3. }
  4. std::cout<<CosineDistance(v1,v2)<<std::endl;

简单的配置一个cmake文件

  1. cmake_minimum_required(VERSION 3.6)
  2. project(tvm_mobilefacenet)
  3. set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -ldl -lpthread")
  4. SET(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
  5. SET(CMAKE_LIBRARY_OUTPUT_DIRECTORY  ${CMAKE_CURRENT_SOURCE_DIR})
  6. SET(HOME_TVM /Users/jackyu/downloads/tvm-0.5)
  7. find_package(OPENCV REQUIRED)
  9. INCLUDE_DIRECTORIES(${OpenCV_INCLUDE_DIRS})
  10. INCLUDE_DIRECTORIES(${HOME_TVM}/include)
  11. INCLUDE_DIRECTORIES(${HOME_TVM}/3rdparty/dmlc-core/include)
  12. INCLUDE_DIRECTORIES(${HOME_TVM}/3rdparty/dlpack/include)
  14. add_executable(tvm_mobilefacenet  tvm_runtime_pack.cc main.cpp)
  15. target_link_libraries(tvm_mobilefacenet    ${OpenCV_LIBS})

Todo：如何在在Android下部署整套人脸识别流程