TVM 简介
TVM 可以将不用的深度学习框架的算法部署到不同的硬件平台来完成计算,网上对于这部分的资料比较少。
我会逐步将自己学习过程中的一些过程记录下来,方便有需求的人。如果能帮助到你,欢迎打赏。
TVM X86 环境配置
代码及编译依赖准备
- sudo apt-get update
- sudo apt-get install -y python python-dev python-setuptools gcc libtinfo-dev zlib1g-dev build-essential cmake
- git clone --recursive https://github.com/dmlc/tvm
- mkdir build
- cp cmake/config.cmake build
- 编辑config.cmake, 来开启或者关闭一些配置。例如要开启cpu 的llvm 编译,首先要获取llvm的安装包: http://releases.llvm.org/download.html
然后将USE_LLVM OFF 改为 set(USE_LLVM /path/to/your/llvm/bin/llvm-config)
编译及环境变量设置
- cd build
cmake ..
make -j4 - 如果一切顺利,则在build目录下会生成:libtvm.so, libtvm_topi.so 。之后可以配置Python环境变量,为了方便以后python代码的更新,建议直接将当前的python环境配置到系统环境中,而不是将python代码拷贝到已经有的环境变量目录。在 ~/.bashrc中添加以下字段:
export TVM_HOME=/path/to/tvm
export PYTHONPATH=TVM_HOME/topi/python:{PYTHONPATH}
然后刷新下环境变量source ~/.bashrc即可 - PC 上最好再安装一些Python的库,方便后面测试以及模型优化:
pip install --user numpy decorator attrs tornado psutil xgboost
测试 Demo
import numpy as np
from tvm import relay
from tvm.relay import testing
import tvm
from tvm.contrib import graph_runtime
batch_size = 1
num_class = 1000
image_shape = (3, 224, 224)
data_shape = (batch_size,) + image_shape
out_shape = (batch_size, num_class)
net, params = relay.testing.resnet.get_workload(
num_layers=18, batch_size=batch_size, image_shape=image_shape)
# set show_meta_data=True if you want to show meta data
print(net.astext(show_meta_data=False))
opt_level = 3
target = tvm.target.cuda()
with relay.build_config(opt_level=opt_level):
graph, lib, params = relay.build_module.build(
net, target, params=params)
# create random input
ctx = tvm.cpu()
data = np.random.uniform(-1, 1, size=data_shape).astype("float32")
# create module
module = graph_runtime.create(graph, lib, ctx)
# set input and parameters
module.set_input("data", data)
module.set_input(**params)
# run
module.run()
# get output
out = module.get_output(0, tvm.nd.empty(out_shape)).asnumpy()
# Print first 10 elements of output
print(out.flatten()[0:10])