TVM 簡介
TVM 可以將不用的深度學習框架的算法部署到不同的硬件平臺來完成計算,網(wǎng)上對于這部分的資料比較少剥哑。
我會逐步將自己學習過程中的一些過程記錄下來师枣,方便有需求的人霞篡。如果能幫助到你血巍,歡迎打賞萧锉。
TVM X86 環(huán)境配置
代碼及編譯依賴準備
- 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, 來開啟或者關(guān)閉一些配置述寡。例如要開啟cpu 的llvm 編譯柿隙,首先要獲取llvm的安裝包: http://releases.llvm.org/download.html
然后將USE_LLVM OFF 改為 set(USE_LLVM /path/to/your/llvm/bin/llvm-config)
編譯及環(huán)境變量設(shè)置
- cd build
cmake ..
make -j4 - 如果一切順利,則在build目錄下會生成:libtvm.so, libtvm_topi.so 鲫凶。之后可以配置Python環(huán)境變量禀崖,為了方便以后python代碼的更新,建議直接將當前的python環(huán)境配置到系統(tǒng)環(huán)境中螟炫,而不是將python代碼拷貝到已經(jīng)有的環(huán)境變量目錄波附。在 ~/.bashrc中添加以下字段:
export TVM_HOME=/path/to/tvm
export PYTHONPATH=TVM_HOME/topi/python:{PYTHONPATH}
然后刷新下環(huán)境變量source ~/.bashrc即可 - PC 上最好再安裝一些Python的庫,方便后面測試以及模型優(yōu)化:
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])