Installing vai_q_pytorch - 2.5 English

Vitis AI User Guide (UG1414)

Document ID
Release Date
2.5 English

vai_q_pytorch has GPU and CPU versions. It supports PyTorch version 1.2~1.10.2 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 encounter 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.10.2)  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
Note: Before running the script, you must check the version of Python, PyTorch, and cuda-toolkit version in the /opt/vitis_ai/scripts/ script and edit them according to your requirement. When choosing PyTorch version and editing the command line, it needs to follow the instructions on 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. 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.10.2) 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 installed PyTorch version 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