首发于转行机器学习的坎坷路
PyTorch模型转TVM模型全流程记录

PyTorch模型转TVM模型全流程记录

概述

PyTorch1.3以后添加了对移动端的支持，我曾尝试过将模型转入移动端，花了很多功夫，把检测+识别的所有代码都转成TorchScript之后，放到移动端运行，却发现在移动端的推理速度比PC慢了好几倍，不得不放弃这个方案。

选择使用TVM主要是看好它能在不同平台进行自动优化的功能，据说这种自动优化的效果能比得上mkl、cudnn之类的加速库。

在这篇文章中，将介绍如何把一个EAST文本检测的PyTorch模型经ONNX转化为TVM模型的过程。

选择这个模型的理由是EAST中包含upsample(interpolate/resize)层，这个算子不算特别常见，在转换过程中，这个upsample层会频繁出错。

希望我的解决过程会给大家提供一点在模型部署过程中解决问题的启发。如果有更好的解决方案，请务必告诉我_(:з」∠)_。

本文主要介绍转换过程中可能遇到的问题以及我的解决方案 代码请参考：github.com/Arctanxy/lea

文章分为四个部分：

  1. 环境准备----编译安装LLVM/TVM/ONNX
  2. PyTorch转ONNX
  3. ONNX转TVM
  4. TVM模型优化

1. 环境准备

计算机软硬件环境：

软件

Ubuntu 18.04 
Python3.7
PyTorch 1.4.0
LLVM/TVM/ONNX均为Github中的master分支最新代码

硬件

CPU: i5-4460 Haswell微架构
内存: 32GB

曾尝试在一台i3-3220（IvyBridge微架构）的虚拟机（2核，4GB内存）上运行，在优化tvm模型后的编译过程中出现KeyError: tile_oh的错误，没能解决就换了机器

这里默认大家都已经安装好PyTorch环境了，没有安装的话可以去PyTorch官网看安装说明。

1.1 LLVM编译

为了编译TVM，首先要编译LLVM，LLVM是一个编译器框架。

tips 1: 我之前尝试过使用预编译的LLVM来编译TVM，但是编译完成之后，TVM还是会提示Target llvm is not enabled. 所以才会使用源码编译
tips 2: 因为库比较大，可以考虑从Gitee下载

下载LLVM

git clone https://gitee.com/mirrors/LLVM.git

编译LLVM

这里选择编译Release版，因为我们也不需要对LLVM进行debug，而且Release版本比较容易编译成功。

mkdir build
cd build
cmake -DCMAKE_BUILD_TYPE=Release ../llvm
make -j4

1.2 TVM编译与PythonAPI安装

编译完LLVM之后就可以编译TVM了，TVM同样需要下载源码进行编译

下载TVM

git clone https://gitee.com/mirrors/tvm.git
git submodule init
git submodule update

编译TVM

mkdir build
cp cmake/config.cmake build

编辑config.cmake 设置set(USE_LLVM /path/to/your/llvm/bin/llvm-config) 当cmake之后看到LLVM的版本信息，说明连接成功

编译

cd build
cmake ..
make -j4
cd ..

编译耗时大约五分钟

编译的最后可能会遇到如下错误

/home/dalalaa/library/tvm/src/target/llvm/llvm_module.cc:242:15: error: no match for ‘operator=’ (operand types are ‘std::__cxx11::string {aka std::__cxx11::basic_string<char>}’ and ‘llvm::StringRef’)

/home/dalalaa/library/tvm/src/target/llvm/codegen_amdgpu.cc:239:39: error: conversion from ‘llvm::StringRef’ to non-scalar type ‘std::__cxx11::string {aka std::__cxx11::basic_string<char>}’ requested

/home/dalalaa/library/tvm/src/target/llvm/codegen_nvptx.cc:218:41: error: conversion from ‘llvm::StringRef’ to non-scalar type ‘std::__cxx11::string {aka std::__cxx11::basic_string<char>}’ requested

通过查询官网文档得知string和llvm::StringRef的转换方法，将llvm_module.cc中242行（codegen_amdgpu.cc和codegen_nvptx.cc做同样的修改）:

target_ = pstr->getString(); 修改为 target_ = pstr->getString().str();

重新编译tvm即可。

安装TVM Python API

我们选择官网提供的第二种安装方式：

cd python; python setup.py install --user; cd ..
cd topi/python; python setup.py install --user; cd ../..

至此，TVM便安装完毕。下面正式开始转换工作。

1.3 ONNX编译安装

本来ONNX直接用pip或者conda安装应该是很方便的，但是ONNX1.4.0老是安装失败；ONNX1.5.0生成的onnx模型无法通过onnx.checker.check_model的检验（可能是Resize算子造成的）；最后，直接使用pip安装的onnx1.6.0在tvm中加载会出现segmentationfault，而从源码安装的ONNX1.6.0就不会。

git clone https://github.com/onnx/onnx.git
sudo apt-get install protobuf-compiler libprotoc-dev
cd ONNX
python setup.py install

2. PyTorch转ONNX

PyTorch转ONNX代码参考:

pytorch.org/docs/stable

在模型的转换过程中的难点在于不常用算子的转换，所以如果选择torchvision中自带的哪几种常见的卷积分类模型的话，这一步就没必要讲了。

我们可以从Github上随便找一个带预训练参数的模型。这里我找来了一个以前在工作中用到过的EAST文本检测模型。

项目地址：

github.com/SakuraRiven/

预训练参数地址：

drive.google.com/file/d

这个项目中用到的了torch.nn.functional.interpolate函数，这个函数不太常用，在转换过程中有不少坑。

tips 3: opset_version选11，interpolate函数转成onnx之后会变成onnx::Resize算子，opset_version11以下都对其支持不完善；
tips 4: 模型转换完成之后无法通过onnx:checker.check_model的检验，可以略过这一步，如果使用onnx1.6.0的话，可以通过这一步检验，但是在tvm加载后又会出现segmentationfault，这也是上面选择源码安装ONNX的原因。

3. ONNX转TVM

ONNX 转 TVM代码参考：

docs.tvm.ai/tutorials/f

转换的过程中出现错误：

Traceback (most recent call last):

  File "east_model.py", line 227, in <module>
    to_tvm()

  File "east_model.py", line 219, in to_tvm
    mod,params = relay.frontend.from_onnx(model,shape_dict)

  File "/home/dalalaa/.local/lib/python3.7/site-packages/tvm-0.7.dev0-py3.7-linux-x86_64.egg/tvm/relay/frontend/onnx.py", line 1863, in from_onnx
    mod, params = g.from_onnx(graph, opset)
... # 错误太长，只贴一部分
...
/home/dalalaa/.local/lib/python3.7/site-packages/tvm-0.7.dev0-py3.7-linux-x86_64.egg/tvm/libtvm.so(dmlc::LogMessageFatal::~LogMessageFatal()+0x32) [0x7fba2ca26342]
  File "/home/dalalaa/library/tvm/src/ir/error.cc", line 133
TVMError: 
Error(s) have occurred. The program has been annotated with them:

最后还有一个红色的粗体字relay.concatenate requires all tensors have the same shape on non-concatenating axes;

看来是在concat的时候尺寸不匹配，而concat的前一步就是upsample（也就是onnx中的resize），可以看到，本来是16x16的特征图变成了128x128：

%70 = image.resize(%69, size=[128, 128], coordinate_transformation_mode="align_corners");

导致尺寸不匹配，无法concat。所以问题还是出在resize算子中。

我们在python3.7/site-packages/tvm-0.7.dev0-py3.7-linux-x86_64.egg/tvm/relay/frontend/onnx.py文件的Resize类中使用print大法，发现错误的原因是pytorch导出的onnx中的resize的参数是（input_size，output_size），而tvm中需要的参数是（input_size, scale）。所以tvm把output_size当成scale来用了，导致转换出错。

这里我尝试过将output_size转换成scale，但是又引发了设计llvm库的错误。

考虑到我的模型中用到的interpolate的scale都是2，我直接在/home/dalalaa/.local/lib/python3.7/site-packages/tvm-0.7.dev0-py3.7-linux-x86_64.egg/tvm/relay/frontend/onnx.py中把resize算子的scale修改成固定值，修改之后，模型转换很顺利，推理验证也没有问题。

显然这不是一个通用的解决方案，如果网络中有多种scale的话，可能需要通过修改onnx节点来解决。如果有大佬了解的的话，麻烦告知。

# scale = infer_value_simulated(inputs[2], params).asnumpy()
scale = np.array([2])

至此便可编译完成，得到east.so / graph.json / params.params三个文件。

测试速度发现：

PyTorch : 0.20s ONNX: 0.16s TVM: 13.1s

TVM的推理速度为何这么慢？一方面是因为我们没用tvm中的runtime模块运行模型，另一方面是因为我们没有对tvm模型进行优化。

使用runtime之后，tvm模型的推理速度变成了1.1s，虽然还是不如有mkl加持的pytorch速度快，但差距不算大。接下来我们将尝试使用autotvm对模型进行优化，看看优化之后的模型速度能否达到pytorch的同等水平。

4. TVM优化

TVM 优化代码参考：

docs.tvm.ai/tutorials/a

优化的过程中可能需要安装依赖库，如果版本不对的话，会有提示，按照提示的版本安装即可：

pip install antlr4-python3-runtime==4.7.2

tips 5:每次开始优化前，最好将east_graph_opt.log文件清空或删除

通过修改early_stopping可以控制搜索次数(调优过程其实有点看运气)

无优化:              speed:1100ms
early_stopping:2    speed:400ms
early_stopping:4    speed:402ms
early_stopping:8    speed:375ms
early_stopping:16   speed:352ms
early_stopping:64   speed:312ms

因为优化时间实在太长了，我没有再继续测试下去了。

模型在i5-4460的电脑上编译很顺利，而在i3-3320的电脑上编译失败，提示KeyError: tile_oh,原因尚未查到，猜测可能跟CPU架构有关。

至此PyTorch到TVM模型的转换工作就全部完成了。

下一步将尝试将模型部署起来。

编辑于 2020-04-19 20:57
PyTorch
TVM
深度学习（Deep Learning）

文章被以下专栏收录