CompressAI Trainer is a training platform that assists in managing experiments for end-to-end neural network-based compression research.

CompressAI Trainer integrates with CompressAI (library), Aim (experiment tracker), Catalyst (training engine), and Hydra (YAML configuration).

CompressAI Trainer integrates with the Aim experiment tracker to display live visualizations of RD curves during training.

Installation

Requirements: Python 3.8+.

First, clone the repositories:

git clone "https://github.com/InterDigitalInc/CompressAI.git" compressai
git clone "https://github.com/InterDigitalInc/CompressAI-Trainer.git" compressai-trainer

Using venv

Create a virtual environment and install as editable:

python3 -m venv venv
source venv/bin/activate
pip install --upgrade pip
pip install --editable ./compressai-trainer
pip install --editable ./compressai

Using poetry

Poetry helps manage version-pinned virtual environments. First, install Poetry:

curl -sSL https://install.python-poetry.org | python3 -

Then, create the virtual environment and install the required Python packages:

cd compressai-trainer

# Install Python packages to new virtual environment.
poetry install
echo "Virtual environment created in $(poetry env list --full-path)"

# Link to local CompressAI source code.
poetry run pip install --editable /path/to/compressai

To activate the virtual environment, run:

poetry shell

Development

To install dependencies for documentation and development, run the following (compatible with both poetry and venv installations):

poetry install --with=dev,docs,tests

You can then build and serve documentation using make:

make docs-serve

Documentation

• Documentation
• Walkthrough (setup, training, visualizations, configuring, customizing, tips and tricks)

Usage

Please see the documentation for a complete walkthrough.

Authors

• Mateen Ulhaq, Fabien Racapé, and InterDigital Emerging Technologies Lab.