源码安装TVM

源码安装 TVM

前言

  • 在实验室 A100 服务器 spack 环境折腾了很久都没有弄好，于是直接转用 docker 在 ubuntu root 环境下安装
  • 水一篇记录一下单纯想放张夏日重现封面图

依赖

  • 以下为最小依赖
apt-get update
apt-get install -y python3 python3-dev python3-setuptools gcc libtinfo-dev zlib1g-dev build-essential cmake libedit-dev libxml2-dev

  • 对于 python 环境，TVM 不支持 python3.9.x 及以上版本，所以这里使用 anconda 虚拟环境，方便版本控制

    • conda 安装略

  • 添加 python 最小依赖

conda install numpy decorator attrs
  • 使用 auto-tuning 的 python 依赖（可选）
conda install tornado psutil 'xgboost<1.6.0' cloudpickle

LLVM

  • TVM 需要使用 LLVM 作为 CPU 生成后端，基本是必装的
apt-get install llvm clang ninja-build
  • 这里用 ninja 作为构建系统

CUDA（可选）

  • TVM 也可以使用 CUDA 作为后端
  • 在 docker 环境中首先要有 CUDA 驱动，docker 运行时需要增加参数--gups
docker run -it --name tvm --gpus all ubuntu
  • 安装 CUDA Toolkit
  • CUDNN（可选）

安装

  • 拉取代码
git clone --recursive https://github.com/apache/tvm tvm
  • 编译准备
cd tvm
mkdir build
cp cmake/config.cmake build
cd build
  • 修改 config.cmake
    • 有非常多的可选项供选择
set(USE_LLVM ON) 
# set(USE_CUDA ON)
# set(USE_CUBLAS ON)
# set(USE_CUDNN ON)

# It is recommended to avoid potential symbol conflicts between different versions LLVM used by TVM and PyTorch. 
# set(USE_LLVM "/path/to/llvm-config --link-static")
# set(HIDE_PRIVATE_SYMBOLS ON)
  • 编译
cmake .. -G Ninja
cmake --build . --parallel 64
  • 安装
cmake --install .
  • 添加环境变量
export TVM_HOME=/path/to/tvm
export PYTHONPATH=$TVM_HOME/python:${PYTHONPATH}
# 或者
# cd python
# python setup.py install

测试

  • gallery 目录下有很多 demo 和教程代码，在满足依赖的情况下可以试一试

  • 以下测试可能都需要安装一些额外的 python 依赖

  • 建议使用多种 conda 环境分别进行测试，避免包冲突

  • onnx 模型测试 TVMC

    • 下载 resnet50 模型
    mkdir myscripts
cd myscripts
wget https://github.com/onnx/models/raw/b9a54e89508f101a1611cd64f4ef56b9cb62c7cf/vision/classification/resnet/model/resnet50-v2-7.onnx
mv resnet50-v2-7.onnx my_model.onnx
    • LLVM
    from tvm.driver import tvmc
model = tvmc.load('my_model.onnx')
package = tvmc.compile(model, target="llvm")
result = tvmc.run(package, device="cpu")
print(result)
    • CUDA
    from tvm.driver import tvmc
model = tvmc.load('my_model.onnx')
package = tvmc.compile(model, target="cuda")
result = tvmc.run(package, device="cuda")
print(result)

  • pytorch 模型测试

    • python gallery/how_to/compile_models/from_pytorch.py

  • autotvm 测试

    • python gallery/how_to/tune_with_autotvm/tune_conv2d_cuda.py

  • autoscheduler 测试

    • python gallery/how_to/tune_with_autoscheduler/tune_conv2d_layer_cuda.py

参考

#记录 #CUDA #conda #docker #TVM #机器学习编译
源码安装TVM
http://example.com/2022/10/04/源码安装TVM/
作者
zty
发布于
2022年10月4日
许可协议