Installing vai_q_pytorch - 3.5 English

Vitis AI User Guide (UG1414)

Document ID
Release Date
3.5 English

vai_q_pytorch has GPU and CPU versions. It supports PyTorch version 1.2~2.0 but does not support PyTorch data parallelism. There are two ways to install vai_q_pytorch:

Install Using Docker Containers

Vitis AI provides a Docker container for quantization tools, including vai_q_pytorch. After running a GPU or CPU container, activate the Vitis AI-PyTorch conda environment,.

conda activate vitis-ai-pytorch
Note: In some cases, if you want to install some packages in the conda environment and encounter issues with permissions, you can create a separate conda environment based on vitis-ai-pytorch instead of using vitis-ai-pytorch directly. The pt_pointpillars_kitti_12000_100_10.8G_1.3 model in AMD Model Zoo is an example of this.

A new conda environment with a specified PyTorch version (1.2~2.0)  can be created using the script. This script does the following:

  • Clones a conda environment from Vitis AI-pytorch.
  • Uninstalls the original PyTorch, torchvision, and vai_q_pytorch packages.
  • Installs the specified version of PyTorch, torchvision.
  • Re-installs vai_q_pytorch from source code.
The following is the command line to create a new conda environment with the script: new_conda_env_name
Note: Before running the script, check the version of Python, PyTorch, and cuda-toolkit version in the script and edit them according to your requirement. When choosing the PyTorch version and editing the command line, follow the instructions on the PyTorch official webpage.

Install from the Source Code

vai_q_pytorch is a Python package designed to work as a PyTorch plugin. It is an open source in Vitis_AI_Quantizer. AMD recommends installing vai_q_pytorch in the conda environment. To do so, follow these steps:

  1. Add the CUDA_HOME environment variable in .bashrc.
    If the CUDA library is installed in ​/usr/local/cuda ​for the GPU version, add the following line into .bashrc. If CUDA is in another directory, change the line accordingly:
    export CUDA_HOME=/usr/local/cuda
    For the CPU version, remove all CUDA_HOME environment variable setting in your .bashrc. It is recommended to clean it up using the following command in the command line of a shell window:
    unset CUDA_HOME
  2. Install PyTorch (1.2~2.0) and torchvision:

    The following code uses PyTorch 1.7.1 and torchvision 0.8.2 for example. You can find detailed instructions for other versions on the PyTorch website.

    pip install torch==1.7.1 torchvision==0.8.2
  3. Install other dependencies:
    pip install -r requirements.txt
  4. Install vai_q_pytorch:
    cd ./pytorch_binding 
    python install
  5. Verify the installation:
    python -c "import pytorch_nndct"
Note: If the installed PyTorch version is 1.4 or higher, import pytorch_nndct before importing the torch in your script. This is caused by a PyTorch bug in versions before 1.4. Refer to PyTorch GitHub issues 28536 and 19668 for details:
import pytorch_nndct
import torch