This repository provides a TensorFlow implementation of LIC-TCM based on:
- Learned Image Compression with Mixed Transformer-CNN Architectures (CVPR 2023 Highlight),
- Channel-wise Autoregressive Entropy Models For Learned Image Compression (ICIP 2020).
Image source: https://arxiv.org/abs/2303.14978.pdf
10/06/2023
- Initial release of this project
This project is based on:
- LIC_TCM, the official PyTorch implementation.
- TensorFlow Compression (TFC), a TF library dedicated to data compression.
- SwinT-ChARM (Transformer-Based Transform Coding, ICLR 2022).
We trained LIC-TCM (large) on the whole ImageNet dataset (not a 300k images subset) using the same configuration described in Section 4.1.1.
We did not apply advanced data pre-processing techniques that are commonly used today - further tuning of hyperparameters and fine-tuning of the model to high-resolution image datasets might provide better results.
We release a single pre-trained model (λ=0.0067) as proof of functional correctness: download.
Currently, we have only tested training/ inference using tf.config.experimental.enable_tensor_float_32_execution(False),
see here for more information. To enable TensorFloat-32 on supported hardware, go to lsk2023.py, l. 51 and set tf.config.experimental.enable_tensor_float_32_execution(True).
We have taken great care to provide exact reimplementations of LIC-TCM (small, medium, large).
| LIC-TCM | Model params (official) | Model params (reimpl) |
|---|---|---|
| large | 76.57 (75.89) | 76.56 |
| medium | 59.13 (58.72) | 59.13 |
| small | 45.18 (44.96) | 45.17 |
The model params (in millions) shown here are re-calculated using Deepspeed Profiler/ get_model_profile(), see here for more information. We generally find that there is a slight discrepancy to the officially reported numbers, which are denoted as (x) for the sake of completeness.
res
├── doc/ # addtional resources
├── eval/ # sample images + reconstructions
├── train_lsk2023/ # model checkpoints + tf.summaries
├── lsk2023/ # saved model
swin-transformers-tf/ # extended swin-transformers-tf implementation
├── changelog.txt # summary of changes made to the original work
├── ...
arch_ops.py # partial translation of compressai.layers, tcm
config.py # model-dependent configurations
lsk2023.py # core of this repo