AIMET Installation in Docker

This page provides instructions to install AIMET package inside a development docker container.

Set variant

Set the <variant_string> to ONE of the following depending on your desired variant
  1. For the PyTorch 2.1 GPU variant, use torch-gpu

  2. For the PyTorch 2.1 CPU variant, use torch-cpu

  3. For the PyTorch 1.13 GPU variant, use torch-gpu-pt113

  4. For the PyTorch 1.13 CPU variant, use torch-cpu-pt113

  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>

Use prebuilt docker image

Follow these instructions to use one of the pre-built docker images:

WORKSPACE="<absolute_path_to_workspace>"
docker_image_name="artifacts.codelinaro.org/codelinaro-aimet/aimet-dev:latest.${AIMET_VARIANT}"
docker_container_name="aimet-dev-<any_name>"

NOTE: Feel free to modify the docker_container_name as needed.

Build docker image locally

Follow these instructions ONLY if you want to build the docker image locally. If not, skip to the next section.

WORKSPACE="<absolute_path_to_workspace>"
docker_image_name="aimet-dev-docker:<any_tag>"
docker_container_name="aimet-dev-<any_name>"
docker build -t ${docker_image_name} -f $WORKSPACE/aimet/Jenkins/Dockerfile.${AIMET_VARIANT} .

NOTE: Feel free to modify the docker_image_name and docker_container_name as needed.

Start docker container

Ensure that a docker named $docker_container_name is not already running; otherwise remove the existing container and then start a new container as follows:

docker ps -a | grep ${docker_container_name} && docker kill ${docker_container_name}

docker run --rm -it -u $(id -u ${USER}):$(id -g ${USER}) \
-v /etc/passwd:/etc/passwd:ro -v /etc/group:/etc/group:ro \
-v ${HOME}:${HOME} -v ${WORKSPACE}:${WORKSPACE} \
-v "/local/mnt/workspace":"/local/mnt/workspace" \
--entrypoint /bin/bash -w ${WORKSPACE} --hostname ${docker_container_name} ${docker_image_name}
NOTE:
  1. Feel free to modify the above docker run command based on the environment and filesystem on your host machine.

  2. If nvidia-docker 2.0 is installed, then add –gpus all to the docker run commands in order to enable GPU access inside the docker container.

  3. If nvidia-docker 1.0 is installed, then replace docker run with nvidia-docker run in order to enable GPU access inside the docker container.

  4. Port forwarding needs to be done in order to run the Visualization APIs from docker container. This can be achieved by running the docker container as follows:

port_id="<any-port-number>"

docker run -p ${port_id}:${port_id} --rm -it -u $(id -u ${USER}):$(id -g ${USER}) \
-v /etc/passwd:/etc/passwd:ro -v /etc/group:/etc/group:ro \
-v ${HOME}:${HOME} -v ${WORKSPACE}:${WORKSPACE} \
-v "/local/mnt/workspace":"/local/mnt/workspace" \
--entrypoint /bin/bash -w ${WORKSPACE} --hostname ${docker_container_name} ${docker_image_name}

Install AIMET packages

From PyPI

Aimet Torch GPU can install from pypi through the following method:

Go to https://pypi.org/project/aimet-torch to identify a version you wish to install

  • For PyTorch 1.13 GPU you should use aimet-torch==1.31.1

  • For Pytorch 2.1.2 GPU you should use aimet-torch >= 1.32.0

sudo apt-get install liblapacke -y
pip install aimet-torch

From Release Package

Alternatively, we host .whl packages for each release at https://github.com/quic/aimet/releases. Identify the release tag of the package you wish to install, then follow the instructions below to install AIMET from the .whl file.

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

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

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

  3. For the PyTorch 1.13 GPU variant, use “torch_gpu-pt113”

  4. For the PyTorch 1.13 CPU variant, use “torch_cpu-pt113”

  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:

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

Install the AIMET packages in the order specified below:

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.10/dist-packages. If that is not the case, please modify it accordingly.

  3. Python dependencies will automatically get installed.

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

Environment setup

Set the common environment variables as follows:

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