MLIC Series [ACMMM 2023 / NCW ICML 2023]

This repo contains the official implementation of MLIC ⁺⁺ .

We highlight MLIC ⁺⁺ , which sloves the quadratic complexity of global context capturing!

MLIC: Multi-Reference Entropy Model for Learned Image Compression [Arxiv] [ACMDL] is accepted at ACMMM 2023 !

MLIC ⁺⁺ : Linear Complexity Multi-Reference Entropy Modeling for Learned Image Compression [Arxiv] [OpenReview] is accepted at ICML 2023 Neural Compression Workshop !

• Compared with version presented at Neural Compression Workshop, ICML 2023 at OpenReview, in the latest arxiv version, we add the details of our prior work presented at ACMMM 2023, new comparisons on complexity and more ablation studies Arxiv.

For the versions of papers, we recommand the latest arxiv version.

Architectures

Performance

Benchmark

Pretrained Models

Update 2024-04-08

I upload the training log when the lambda is 0.0250. The model is trained on 4 GPU cards with ddp support.

I fix the LatentResidualPrediction and SynthesisTransform, you should use LatentResidualPrediction and SynthesisTransform instead of LatentResidualPredictionOld and SynthesisTransformOld. The parameter number of MLIC ⁺⁺ becomes 83.5M. The modification leads to no performance drop.

Update checkpoint: https://disk.pku.edu.cn/link/AABED8912D2502477EB37C18FC7F2B2612

code: ujrv

Google Drive: https://drive.google.com/file/d/1FWPezuHLTQhDmEhShViI3XOSXA5u_Bya/view?usp=sharing

Bug Fixes (2023-09)

I fix the implementation of LinearGlobalIntraContext.

Old pre-trained models are put below. The performance is slightly better than before on Kodak and Tecnick.

Old Weights (2023-09)

To use old weights, you should use LatentResidualPredictionOld and SynthesisTransformOld.

The structure of the provided weights is

{
      "epoch": epoch + 1,
      "state_dict": net.state_dict(),
      "loss": loss,
      "optimizer": optimizer.state_dict(),
      "aux_optimizer": aux_optimizer.state_dict(),
      "lr_scheduler": lr_scheduler.state_dict(),
}

Training

Settings

We train each model on a single Tesla A100 GPU. The batch size is set to $32$. The initial patch size is set to $256\times 256$. We set the patch size to $512\times 512$ after $1.2$ M steps.

Training Set

Training list is provided. These images are from DIV2K, Flicker2K, CLIC Train, COCO, ImageNet. Most JPG images are downsampled and the downsampled images are stored in PNG format. We use following function from PIL to downsample images.

img.resize((new_width, new_height), Image.ANTIALIAS)

Command

Example command is provided Here.

Testing

Example command is provided Here.

Environment

CompressAI 1.2.0b3

Pytorch 2.0.1

Contact

If you have any questions about MLIC, please contact Wei Jiang ( [email protected] or [email protected] )

Citation

If you find our papers and this repo useful, kindly cite:

MLIC

@inproceedings{jiang2023mlic,
  title={MLIC: Multi-Reference Entropy Model for Learned Image Compression},
  author={Jiang, Wei and Yang, Jiayu and Zhai, Yongqi and Ning, Peirong and Gao, Feng and Wang, Ronggang},
  doi = {10.1145/3581783.3611694},
  booktitle={Proceedings of the 31st ACM International Conference on Multimedia},
  pages={7618--7627},
  year={2023}
}

MLIC ⁺⁺

@inproceedings{jiang2023mlicpp,
  title={MLIC++: Linear Complexity Multi-Reference Entropy Modeling for Learned Image Compression},
  author={Jiang, Wei and Wang, Ronggang},
  booktitle={ICML 2023 Workshop Neural Compression: From Information Theory to Applications},
  year={2023},
  url={https://openreview.net/forum?id=hxIpcSoz2t}
}