Awesome
Finite-State Autoregressive Entropy Coding for Efficient Learned Lossless Compression
This is the official Pytorch implementation of our ICLR paper:
Finite-State Autoregressive Entropy Coding for Efficient Learned Lossless Compression
<img src="imgs/main.png" width="800">Finite-State Autoregressive Entropy Coding is a VAE-based compression method designed for better compression ratio and computational efficiency. It extends Asymmetric Numeral Systems (ANS) with a lookup-table-based autoregressive model, which efficiently performs autoregressive encoding/decoding that improves compression ratio, even without parallel computation. Besides, the Straight-Through Hardmax Quantization (STHQ) is proposed to enhance the optimization of discrete latent space in VAE.
Setup
Hardware Requirements
- CUDA compatible GPU (Optional but highly recommended for training. For testing, no GPU is OK)
- Memory > 16G * GPUNum (Optional but required for training on ImageNet64. We force full dataset caching for ImageNet64, which requires ~12G memory per dataset instance.)
Software Requirements
- Linux/Mac system (Windows may be possible but not tested)
- gcc&g++>=7 (Required, need c++17 support to compile)
- python>=3.7 (Recommended version, need pytorch-lightning support)
- pytorch>=1.7 (Recommended version, need pytorch-lightning support)
- cudatoolkit>=10.2 (Optional but highly recommended, lower versions may be functional)
- craystack (Optional, if you want to compare with BB-ANS)
The recommended environment setup script with conda:
conda create -n cbench python=3.7
conda activate cbench
conda install -c pytorch pytorch==1.7.1 torchvision==0.8.2 cudatoolkit=10.2
# conda install -c pytorch pytorch==1.12.1 torchvision==0.13.1 cudatoolkit=11.3
# if gcc version < 7
conda install -c conda-forge gcc gxx
pip install -r requirements.txt
python setup.py build develop
(Optional) Environment setup according to your machine
See configs/env.py
(Optional) 3rdparty setup
git submodule update --init --recursive
cd 3rdparty/craystack
python setup.py build develop
Dataset Prepare
We use 5 datasets in our experiments:
- CIFAR10 : No need for preparation, torchvision would automatically download for you.
- ImageNet32/64 : Download downsampled images 32x32 and 64x64 and decompress them into data/ImageNet. The result folders would be:
- data/ImageNet/Imagenet32_train_npz
- data/ImageNet/Imagenet32_val_npz
- data/ImageNet/Imagenet64_train_npz
- data/ImageNet/Imagenet64_val_npz
- OpenImages : We follow L3C to preprocess OpenImages as training/validation set. The result folders would be:
- data/openimages/clean_train
- data/openimages/clean_val
- Kodak : Download all images to data/Kodak
- CLIC 2020 validation : Download and unzip all images to data/CLIC/valid
Code Structure
- configs : All parameter setting configs
- cbench : All algorithm codes
- benchmark : Tools for experiments and comparisons
- csrc : c/c++ codes, with an efficient ANS-based entropy coder interface/implementation. (Should be more efficient than CompressAI implementation, we use pybind11 Tensor support to avoid memory format conversion between python and c++.)
- codecs : All compression codecs
- data : Tools for data loading and processing
- modules : Mainly compression modules
- prior_model/prior_coder/latent.py : Regarding the proposed FSAR, please refer to CategoricalAutoregressivePriorDistributionPriorCoder; For STHQ, please refer to StochasticVQAutoregressivePriorDistributionPriorCoder.
- nn : Neural network codes.
- utils : Misc tools.
- tests : To test if some modules are functional.
- tools : Tools to run experiments.
- run_benchmark.py : Use this for all experiments!
Experiments
For any experiment you want to run (including Training/Validation/Testing, thanks to pytorch-lightning and our BasicLosslessCompressionBenchmark):
python tools/run_benchmark.py [config_file]
You can use tensorboard to visualize the training process.
tensorboard --logdir experiments
Experiment List
NOTE: If an GPU out-of-memory error occured, adjust batch_size_total in each config file. In our experiments, we use 8 A100 for most configs so the batch_size_total might be too large.
- configs/coding_exp_cf10.py (and im32, im64, oi) : Main codec comparison experiments. Could test compression ratio and speed for different image codecs.
- configs/model_shallow_exp_cifar10.py (and im32, im64) : Main latent model comparison experiments. Could test compression ratio and speed for different latent models.
- configs/model_shallow_exp_abl_cifar10.py : Ablation study for different VQ methods.
- configs/model_abl_bb_exp_cifar10.py : Ablation study for different backbone setting.
- configs/dsar_exp.py : Ablation study for data/observation space autoregressive coding.
Model Implementation List
See configs/presets.
Our full model is tagged as "V2DVQ-c2-FSAR-O2S-catreduce1.5"
Pretrained Models
TBA
tools/run_benchmark.py can automatically look for config.pth in a given directory to build the benchmark. Therefore, to test a pretrained model, simply run:
python tools/run_benchmark.py [model_directory]
Citation
@inproceedings{
zhang2024finitestate,
title={Finite-State Autoregressive Entropy Coding for Efficient Learned Lossless Compression},
author={Yufeng Zhang and Hang Yu and Jianguo Li and Weiyao Lin},
booktitle={The Twelfth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=D5mJSNtUtv}
}
Contact
TBA