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

	Kodak	Tecnick	CLIC Pro Valid
VTM-17.0 Intra	0.00	0.00	0.00
STF (CVPR'22)	-2.48	-2.75	+0.42
WACNN (CVPR'22)	-2.95	-	+0.04
ELIC (CVPR'22)	-5.95	-	-
LIC-TCM Large (CVPR'23)	-10.14	-11.47	-8.04
MLIC (ACMMM'23)	-8.05	-12.73	-8.79
MLIC+ (ACMMM'23)	-11.39	-16.38	-12.56
MLIC++ (NCW ICML'23)	-13.39	-17.59	-13.08

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 LatentResidualPredictionFix and SynthesisTransformFix. 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

Bug Fixes

I fix the implementation of LinearGlobalIntraContext.

New pre-trained models are available now. The performance is slightly better than before on Kodak and Tecnick.

New Weights

Lambda	Metric	Link	Lambda	Metric	Link
0.0018	MSE	PKUDisk, GoogleDrive	2.4	MS-SSIM	PKUDisk, GoogleDrive
0.0035	MSE	PKUDisk, GoogleDrive	4.58	MS-SSIM	PKUDisk, GoogleDrive
0.0067	MSE	PKUDisk, GoogleDrive	8.73	MS-SSIM	PKUDisk, GoogleDrive
0.0130	MSE	PKUDisk, GoogleDrive	16.64	MS-SSIM	PKUDisk, GoogleDrive
0.0250	MSE	PKUDisk, GoogleDrive	31.73	MS-SSIM	PKUDisk, GoogleDrive
0.0483	MSE	PKUDisk, GoogleDrive	60.5	MS-SSIM	PKUDisk, GoogleDrive

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 ( wei.jiang1999@outlook.com or jiangwei@stu.pku.edu.cn )

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}
}