Ctrl+K

TVM 准备

TVM 准备#

安装运行环境#

本项目使用 Anaconda3 配置 Python 环境，还需要配置一些包：

注意

由于 TVM 高度依赖于原生系统的可用库，推荐从其源码安装它。详细内容可参考：从源码安装 TVM。

本书需要一些 TVM 编译配置：

set(USE_LLVM ON)

如果您计划在 Nvidia GPU 上运行，那么您也需要

set(USE_CUDA ON)

也不要忘记启用 cython，它可以加速性能。您只需要在 TVM 源文件夹中运行 make cython。

（可选）安装 MXNet GPU 版本（[Chen et al., 2015]）可以：

pip install mxnet-cu116

（可选）安装 PyTorch 可以：

conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

（可选）安装 ONNX，可以：

pip install onnx onnxoptimizer

便捷安装#

注意

当前仅仅支持 Unubtu 平台。

可以直接利用脚本 task.py （需要安装 d2py 和 invoke）编译 TVM。

  1. 构建 build/ 目录，且创建并配置 build/cmake/config.cmake 文件。

    invoke config

    若要启用 CUDA，可以：

    invoke config --cuda

  2. 构建并编译 TVM：

    invoke make

  3. 若要生成 TVM 文档，可：

    invoke update
invoke doc

备注

如若想更新 TVM 文档翻译只需要

invoke doc.update
invoke doc.update -l zh_CN

  1. 如果想要把 tvm 和 vta 作为全局环境的包，可以在 tvm 项目下创建软链接：

    invoke ln-env --root tvm项目根目录 --target python环境路径 --python-version 3.10

  2. 如果 TVM 和 VTA 仅仅是作为局部使用，可以使用 pdm 进行管理。

关于 tvm_book 库#

为了便于本项目程序的复用，定制了 tvm_book 库，可以直接在项目根目录直接安装：

pip install .

如若想要本地化项目文档，可以

pip install .[doc]

如若不想将 tvm_book 安装到全局环境，可以在本次运行代码前添加如下函数用于设置 Python 临时环境：

def set_env(num, current_path='.'):
    '''
    num 表示相对于 current_path 的父级根目录级别
    '''
    import sys
    from pathlib import Path

    ROOT = Path(current_path).resolve().parents[num]
    sys.path.extend([str(ROOT/'src')]) # 设置 `tvm_book` 环境
    from tvm_book.config.env import set_tvm
    # 设置 TVM 环境
    set_tvm(TVM_ROOT)

备注

set_tvm 需要自行配置以适配您的设备。

安装的 FAQs#

安装 MXNet GPU 版本可能需要配置 NCCL，详细内容见：安装 NCCL。

备注

这里是 Ubuntu22.04，可以替换为其他版本。

  1. 安装 CUDA：

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub
sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/ /"
sudo apt-get update
sudo apt-get -y install cuda

  1. 安装 NCCL：

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb
sudo dpkg -i cuda-keyring_1.0-1_all.deb
sudo apt-get update

同时还需要：

sudo apt install libnccl2=2.13.4-1+cuda11.7 libnccl-dev=2.13.4-1+cuda11.7

TVM 安装可以是：

pip install mlc-ai-nightly -f https://mlc.ai/wheels

pip install mlc-ai-nightly-cu110 -f https://mlc.ai/wheels

自动调优#

为了在 TVM 中使用 autotvm 包，需要安装一些额外的依赖项。

pip install psutil cloudpickle

还需要安装 xgboost：

pip install xgboost 

conda install -c conda-forge py-xgboost-gpu

查看模型#

查看模型的 input/shape_dict 的推荐方法是通过 netron。打开模型后，单击第一个节点，在 inputs 部分查看名称和形状。

其他#

加速 scikit-learn 可以安装 scikit-learn-intelex：

conda install scikit-learn-intelex
目录