CompressLab

A PyTorch-based framework for deep-learning-based data compression research.

Features

• All implementations of models and datasets are modularized and registered in a registry. The training and testing pipelines can be easily configured by YAML files. Please refer to the example configuration files in the config folder and see the details of available options in compresslab/utils/config.py.
• Add !include tag to include other YAML files, and !loop tag to generate repeated structures in the configuration files. Please refer to the examples in config folder.
• For the vae-based image compression models compresslab/nn/lossy_image_compression, training with several rate-distortion points simultaneously is supported in one run.
• The framework is built on PyTorch Lightning, which makes it easy to run in DDP mode and adjust the precision. Moreover, it supports logging via TensorBoard, and resume from the checkpoint automatically.
• A general metrics computation module is provided, and you can easily obtain the final results in the output folder after testing.

Installation

conda create -n compresslab python=3.9
conda activate compresslab
pip install -e .

How to use

Run

python run.py --list

to check all registered modules.

Run

python run.py --config config/your_config.yaml

to train a model with the specified configuration file.

Run

tensorboard --logdir output

to monitor the training process.

Dataset preparation

download.py can be used to download and extract the datasets automatically. For example, to download the CLIC2020 Professional Validation dataset, run

python dataset/download.py --dataset CLIC2020_professional_val

Current supported datasets:

• Vimeo90k
• CLIC2020
• DIV2K
• Tecnick
• UVG
• Kodak (Already included in the repository)

Run

python dataset/download.py --list

to check all supported datasets.

Available models

VAE-based lossy image compression

• All models implemented in CompressAI:
  • FactorizedPrior
  • ScaleHyperprior
  • MeanScaleHyperprior, JointAutoregressiveandHierarchicalPriors
  • Cheng2020
  • Cheng2020Checkerboard
  • ELIC
• ChARM
• TCM
• MLIC

Generative image compression (Only support evaluation of the pre-trained models now)

• DiffEIC
• StableCodec

TODO

• Refactor the code of video compression models.
• Load video dataset directly from .yuv files.
• Try to implement finetining from checkpoints or two-stage training in one run.
• Support automatically padding the input image to the required size.
• benchmark.py can't give the correct bd-rate curve. Fix it later.

Related links

• PyTorch Lightning
• CompressAI
• ChARM-PyTorch
• LIC-TCM
• MLIC
• DVC-PyTorch
• PyTorchVideoCompression
• DiffEIC
• StableCodec
• CLIC
• UVG
• Kodak