Installing vai_q_pytorch - 2.0 English

Vitis AI User Guide (UG1414)

Document ID
Release Date
2.0 English

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

Install Using Docker Containers

The Vitis AI provides a Docker container for quantization tools, including vai_q_pytorch. After running a GPU/CPU container, activate the Conda environment, vitis-ai-pytorch.

conda activate vitis-ai-pytorch
Note: In some cases, if you want to install some packages in the conda environment and meet permission problems, 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 Xilinx Model Zoo is an example of this.
A new Conda environment with a specified PyTorch version (1.2~1.9)  can be created using the /opt/vitis_ai/scripts/ script. This script clones a Conda environment from vitis-ai-pytorch, uninstalls the original PyTorch, Torchvision and vai_q_pytorch packages, and then installs the specified version of PyTorch, Torchvision, and re-installs vai_q_pytorch from source code. The following is the command line to create a new Conda environment with the script:
/opt/vitis_ai/scripts/ new_conda_env_name

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. It is recommended to install vai_q_pytorch in the Conda environment. To do so, follow these steps:

  1. Add the CUDA_HOME environment variable in .bashrc.
    For the GPU version, if the CUDA library is installed in /usr/local/cuda, add the following line into .bashrc. If CUDA is in other 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 cleanup it in command line of a shell window by running the following command:
    unset CUDA_HOME
  2. Install PyTorch (1.2-1.9) and Torchvision.

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

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