端到端图像编码代码复现

2025.8.3

1 基本环境

使用bitahub(https://bitahub.ustc.edu.cn/)集群。

使用PyTorch 1.12，Python3.8环境（bitahub基础镜像10.11.3.8:5000/bitahub/deepo:py38_cu113）。

额外库有pytorch_msssim einops compressai==1.1.8 opencv-python h5py tensorboard，使用如下命令安装：

pip install pytorch_msssim einops compressai==1.1.8 opencv-python h5py tensorboard -i https://pypi.tuna.tsinghua.edu.cn/simple

2 代码使用

模型文件在/models/目录下，共有五个基本模型，其中如下四个模型的代码复现工作已经完成：

from .factorized import Factorized
from .hyperprior import Hyperprior
from .joint_autogressive import JointAutoregressive
from .channel_ar import Channel_AR

2.1 模型训练

使用train.py进行训练，训练指令如下，其中可用模型名在/zoo/__init__.py中。

lambda代表训练损失中L=R+λD中的拉格朗日乘子λ，用于控制码率，训练模型时，分别取λ=[0.0018, 0.0054, 0.0162, 0.0483]训练4个码率点。

python train.py --model [模型名] --lambda [lambda值]

2.2 模型测试

模型训练好后，使用test.py进行测试，测试指令如下，其中可用模型名在/zoo/__init__.py中，模型路径为训练好的、需要进行测试的模型的路径。

python test.py --model_name [模型名] --path [模型路径]

3 代码复现

3.1 Factorized模型

论文URL：https://arxiv.org/pdf/1611.01704

测试结果：

BPP	PSNR	MS-SSIM
0.13272	27.13278	0.90851
0.29241	29.75045	0.94911
0.58698	32.75182	0.97492
1.06703	36.05223	0.98821

与官方测试结果对比：

BD-PSNR 0.10527 BD-RATE -2.4561

3.2 Hyperprior模型

论文URL：https://arxiv.org/pdf/1802.01436

测试结果：

BPP	PSNR	MS-SSIM
0.13769	27.56897	0.91607
0.29632	30.4383	0.95388
0.56174	33.67108	0.97761
0.97883	37.00354	0.9892

与官方测试结果对比：

BD-PSNR -0.45007 BD-RATE 9.93335

3.3 JointAutoregressive模型

论文URL：https://arxiv.org/pdf/1809.02736

测试结果：

BPP	PSNR	MS-SSIM
0.11721	27.9865	0.91865
0.25955	30.83388	0.95632
0.5144	33.85291	0.97779
0.91596	37.05343	0.989

与官方测试结果对比：

BD-PSNR -0.10707 BD-RATE 2.54386

3.4 Channel_AR模型

论文URL：https://arxiv.org/pdf/2007.08739

测试结果：

BPP	PSNR	MS-SSIM
0.13361	28.78062	0.9283
0.28637	31.4402	0.96203
0.55066	34.39043	0.98094
0.98743	37.60767	0.99042

与官方测试结果对比：

BD-PSNR 0.33401 BD-RATE -7.17225

4 附录