AIMET Installation and Setup

This page provides instructions to install AIMET package on Ubuntu 20.04 LTS with Nvidia GPU. Please follow the instructions in the order provided, unless specified otherwise.

NOTE:

  1. Please pre-pend the “apt-get install” and “pip3 install” commands with “sudo -H” as appropriate.

  2. These instructions assume that pip packages will be installed in the path: /usr/local/lib/python3.8/dist-packages. If that is not the case, please modify it accordingly.

Install prerequisite packages

Install the basic pre-requisite packages as follows:

apt-get update
apt-get install python3.8 python3.8-dev python3-pip
python3 -m pip install --upgrade pip
apt-get install --assume-yes wget gnupg2

If you have multiple python versions installed, set the default python version as follows:

update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.8 1
update-alternatives --set python3 /usr/bin/python3.8

Install GPU packages

NOTE:

  1. Do this section ONLY for the GPU variants.

  2. The released AIMET GPU packages were tested with the following CUDA toolkit versions:

    1. PyTorch 1.9 GPU variant: CUDA Toolkit 11.1.1

    2. PyTorch 1.13 GPU variant: CUDA Toolkit 11.6.2

    3. TensorFlow GPU variant: CUDA Toolkit 11.2.2

    4. ONNX GPU variant: CUDA Toolkit 11.1.1

  3. The instructions in the sub-sections below correspond to our tested versions above. Visit this page https://developer.nvidia.com/cuda-toolkit-archive to obtain the correct version of the CUDA toolkit for your environment.

Install GPU packages for PyTorch 1.9 or ONNX

NOTE:

  1. Do this section ONLY for the PyTorch or ONNX GPU variants.

  2. Visit this page https://developer.nvidia.com/cuda-11.1.1-download-archive to obtain the exact and up-to-date installation instructions for your environment.

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.1.1/local_installers/cuda-repo-ubuntu2004-11-1-local_11.1.1-455.32.00-1_amd64.deb
dpkg -i cuda-repo-ubuntu2004-11-1-local_11.1.1-455.32.00-1_amd64.deb
apt-key add /var/cuda-repo-ubuntu2004-11-1-local/7fa2af80.pub
echo "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64 /" > /etc/apt/sources.list.d/cuda.list
echo "deb https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu2004/x86_64 /" > /etc/apt/sources.list.d/nvidia-ml.list
apt-get update
apt-get -y install cuda

wget http://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu2004/x86_64/nvidia-machine-learning-repo-ubuntu2004_1.0.0-1_amd64.deb
dpkg -i nvidia-machine-learning-repo-ubuntu2004_1.0.0-1_amd64.deb
apt-get update

Install GPU packages for PyTorch 1.13

NOTE:

  1. Do this section ONLY for the PyTorch 1.13 variant.

  2. Visit this page https://developer.nvidia.com/cuda-11-6-2-download-archive to obtain the exact and up-to-date installation instructions for your environment.

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.6.2/local_installers/cuda-repo-ubuntu2004-11-6-local_11.6.2-510.47.03-1_amd64.deb
dpkg -i cuda-repo-ubuntu2004-11-6-local_11.6.2-510.47.03-1_amd64.deb
apt-key add /var/cuda-repo-ubuntu2004-11-6-local/7fa2af80.pub
echo "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64 /" > /etc/apt/sources.list.d/cuda.list
echo "deb https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu2004/x86_64 /" > /etc/apt/sources.list.d/nvidia-ml.list
apt-get update
apt-get -y install cuda

wget http://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu2004/x86_64/nvidia-machine-learning-repo-ubuntu2004_1.0.0-1_amd64.deb
dpkg -i nvidia-machine-learning-repo-ubuntu2004_1.0.0-1_amd64.deb
apt-get update

Install GPU packages for TensorFlow

NOTE:

  1. Do this section ONLY for the TensorFlow GPU variant.

  2. Visit this page https://developer.nvidia.com/cuda-11.2.2-download-archive to obtain the exact and up-to-date installation instructions for your environment.

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.2.2/local_installers/cuda-repo-ubuntu2004-11-2-local_11.2.2-460.32.03-1_amd64.deb
dpkg -i cuda-repo-ubuntu2004-11-2-local_11.2.2-460.32.03-1_amd64.deb
apt-key add /var/cuda-repo-ubuntu2004-11-2-local/7fa2af80.pub
echo "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64 /" > /etc/apt/sources.list.d/cuda.list
echo "deb https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu2004/x86_64 /" > /etc/apt/sources.list.d/nvidia-ml.list
apt-get update
apt-get -y install cuda

wget http://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu2004/x86_64/nvidia-machine-learning-repo-ubuntu2004_1.0.0-1_amd64.deb
dpkg -i nvidia-machine-learning-repo-ubuntu2004_1.0.0-1_amd64.deb
apt-get update

Install AIMET packages

Go to https://github.com/quic/aimet/releases and identify the release tag of the package you want to install.

Set the <variant_string> to ONE of the following depending on your desired variant

  1. For the PyTorch 1.13 GPU variant, use “torch_gpu”

  2. For the PyTorch 1.13 CPU variant, use “torch_cpu”

  3. For the PyTorch 1.9 GPU variant, use “torch_gpu_pt19”

  4. For the PyTorch 1.9 CPU variant, use “torch_cpu_pt19”

  5. For the TensorFlow GPU variant, use “tf_gpu”

  6. For the TensorFlow CPU variant, use “tf_cpu”

  7. For the ONNX GPU variant, use “onnx_gpu”

  8. For the ONNX CPU variant, use “onnx_cpu”

export AIMET_VARIANT=<variant_string>

Replace <release_tag> in the steps below with the appropriate tag:

export release_tag=<release_tag>

Set the package download URL as follows:

export download_url="https://github.com/quic/aimet/releases/download/${release_tag}"

Set the common suffix for the package files as follows:

NOTE: Set wheel_file_suffix to cp38-cp38-linux_x86_64.whl OR cp36-cp36m-linux_x86_64 OR cp37-cp37m-linux_x86_64 OR py3-none-any as appropriate depending on the actual wheel filename(s) on the https://github.com/quic/aimet/releases.

export wheel_file_suffix="cp38-cp38-linux_x86_64.whl"

Install the AIMET packages in the order specified below:

NOTE: Python dependencies will automatically get installed.

python3 -m pip install ${download_url}/AimetCommon-${AIMET_VARIANT}_${release_tag}-${wheel_file_suffix}

# Install ONE of the following depending on the variant
python3 -m pip install ${download_url}/AimetTorch-${AIMET_VARIANT}_${release_tag}-${wheel_file_suffix} -f https://download.pytorch.org/whl/torch_stable.html
# OR
python3 -m pip install ${download_url}/AimetTensorflow-${AIMET_VARIANT}_${release_tag}-${wheel_file_suffix}
# OR
python3 -m pip install ${download_url}/AimetOnnx-${AIMET_VARIANT}_${release_tag}-${wheel_file_suffix}

python3 -m pip install ${download_url}/Aimet-${AIMET_VARIANT}_${release_tag}-${wheel_file_suffix}

Install common debian packages

Install the common debian packages as follows:

cat /usr/local/lib/python3.8/dist-packages/aimet_common/bin/reqs_deb_common.txt | xargs apt-get --assume-yes install

NOTE: Do the following ONLY for the PyTorch variant packages.

cat /usr/local/lib/python3.8/dist-packages/aimet_onnx/bin/reqs_deb_torch_common.txt | xargs apt-get --assume-yes install

NOTE: Do the following ONLY for the ONNX variant packages.

cat /usr/local/lib/python3.8/dist-packages/aimet_onnx/bin/reqs_deb_onnx_common.txt | xargs apt-get --assume-yes install

Install tensorflow GPU debian packages

NOTE: Do this ONLY for the TensorFlow GPU package.

cat /usr/local/lib/python3.8/dist-packages/aimet_tensorflow/bin/reqs_deb_tf_gpu.txt | xargs apt-get --assume-yes install

Install torch GPU debian packages

NOTE: Do this ONLY for the PyTorch GPU package.

cat /usr/local/lib/python3.8/dist-packages/aimet_torch/bin/reqs_deb_torch_gpu.txt | xargs apt-get --assume-yes install

Install ONNX GPU debian packages

NOTE: Do this ONLY for the ONNX GPU package.

cat /usr/local/lib/python3.8/dist-packages/aimet_onnx/bin/reqs_deb_onnx_gpu.txt | xargs apt-get --assume-yes install

Replace Pillow with Pillow-SIMD

Optional: Replace the Pillow package with Pillow-SIMD as follows:

python3 -m pip uninstall -y pillow
python3 -m pip install --no-cache-dir Pillow-SIMD==9.0.0.post1

Replace onnxruntime with onnxruntime-gpu

NOTE: Do this ONLY for the PyTorch GPU package.

export ONNXRUNTIME_VER=$(python3 -c 'import onnxruntime; print(onnxruntime.__version__)')
python3 -m pip uninstall -y onnxruntime
python3 -m pip install --no-cache-dir onnxruntime-gpu==$ONNXRUNTIME_VER

Post installation steps

ln -s /usr/lib/x86_64-linux-gnu/libjpeg.so /usr/lib

NOTE: Do the following step ONLY for the PyTorch or Tensorflow GPU packages.

# NOTE: Please modify the below command depending on the version of your CUDA driver toolkit
ln -s /usr/local/cuda-11.1 /usr/local/cuda

Environment setup

Set the common environment variables as follows:

source /usr/local/lib/python3.8/dist-packages/aimet_common/bin/envsetup.sh