Transformer-based Transform Coding [ICLR 2022] in CompressAI + Pytorch Lightning

[ali-zafari]

Hi everyone!

I have integrated the four models described in the recent ICLR paper Transformer-based Transform Coding into the CompressAI framework. You can find my implementation here:

github.com/ali-zafari/TBTC.

I tried to follow the code structure of CompressAI to make it easily accessible for anyone familiar with this great library in PyTorch. This TensorFlow implementation (SwinT-ChARM) is used as reference.

Models are defined in TBTC/compressai/models/qualcomm.py:

• Conv-Hyperprior: Convolutional analysis/synthesis transforms + hyperprior entropy model
• Conv-ChARM: Convolutional analysis/synthesis transforms + channel-wise autoregressive entropy model
• SwinT-Hyperprior: Shifted window transformer-based analysis/synthesis transforms + hyperprior entropy model
• SwinT-ChARM: Shifted window transformer-based analysis/synthesis transforms + channel-wise autoregressive entropy model

To do the training, I wrapped CompressAI-based models in Lightning module to make the multi-gpu training and logging/checkpointing easier. You can also download a sample checkpoint for each of the models to verify their performance with the results reported in the paper.

Hope it would be useful.