About prob_model and training

[TNTWEN]

你好！我在尝试您的项目时有几点疑问想要请教。
1、论文中写到超先验部分使用没有任何假设分布的channel-wise fully factorized density model。但在代码中的prob_model似乎是channel-wise 的高斯分布。 使用高斯分布有何优势？
2、我将entroformer 迁移到compressai框架下实验，训练集采用openimage的train_f并做了筛选，大约有9W张图片，以（384，384）的patchsize训练。网络输入范围为[0,1],超先验部分采用了 fully factorized density model，训练策略与您的实现相同。 加载了您的预训练模型进行训练，但效果不太理想，在Kodak上测试结果和Minnen (2018)差不多。 但相同的数据集在您的代码下训练效果就很好。想听听您的意见，不知道是不是fully factorized density model导致的

非常感谢您的工作！

[mx54039q]

我已收到你的邮件。                                               by    钱一琛

[mx54039q]

你好,

1. 这里的描述是勘误, 确实使用的是channel-wise的高斯分布, 感谢指出, 后续会修正. 高斯分布会比较贴切数据的原始分布, 在效果上应该会比fully factorized density model更好 (经验性的), 但对整体模型的性能影响不大.
2. 假设训练策略和其他的结构都没问题的话, 我能想到的可能的原因有两点: 1). 加载我的预训练模型之后, fully factorized density model是重新初始化的, 可能在训练初始的时候会使得context model和hyperprior model失衡, 这里你可以check一下bpp_z是否正常. 2). 图片输入的范围 [-1,1] -> [0,1], 这一块可能会影响到Transformer结构的优化.

[TNTWEN]

你好！感谢您的回复！
确实有可能是加载预训练权重的问题，我也正在尝试从头开始训练，看看是否会有提升。

非常感谢您的工作！祝好！

[TNTWEN]

@mx54039q
你好！阅读代码时发现该处self.dim是不是应该改成self.key_value_proj_dim，即多头注意力的dk取64而非384

entroformer/module/entroformer_helper.py

Line 115 in 1c359bf

self.scale = self.dim ** -0.5 if config.scale else 1.

[mx54039q]

你说的没错, 这一块看来是我的实现错误, 不过训练影响不大. 你可以修改一下, 如果有正向提升的话麻烦告知我一下, 谢谢. :)

[TNTWEN]

@mx54039q
我已经找到了之前训练效果不佳的原因，之前训练时我对于y只用了add noise的操作，没有用ste。 由于我只用bidirectional，于是我把forward过程做到和compress，decompress一致，无需add noise操作。 在compressai的编码方案中，我使用了quantize_ste(y - means_hat) + means_hat。

我将图像输入域改为了[-1,1]，筛选了1.5W张图片作为数据集，加载了您的权重，目前训练了两个模型，效果还不错。在Kodak上compress/decompress的结果如下：
λ=0.02 bpp：0.607 PSNR：35.275
λ=0.18 bpp：1.623 PSNR： 41.401
应该与您paper的结果很相近了。

我现在正在尝试您的random mask的pretrain方案。我修改了dk进行pretrain，然后看看finetune的效果如何，再与您交流^-^