Learned image compression

作者，你好！你的代码对我很有帮助，但是你的训练数据实在太大，国内用户下载google drive太慢。所以我把你的dataset里面注释那行transforms.RandomResizedCrop(self.image_size)加上了，使用了256的裁剪。但是我训练数据使用121326张图片，大概迭代1000次（还在第0个epoch），就报pillow太大的错误了ValueError: Decompressed Data Too Large。但是将图片训练集改到856张的训练集后，就没有出现问题。

Hello

Just wanted to know whether this implementation is a simple autoencoder based one or not? As compressed_feature_renorm is directly fed into the encoder. We are not calculating any mean or standard deviation here.

Dear Jiaheng Liu,

I am keep trying to reproduce result but number of iterations are too big. Thus, It takes long time on moderate GPU. I am requesting you to please give me trained model.

Kind Regards, Khawar

Hello. I'm getting a lot of help in studying.

First of all, thank you very much.

I'm writing this because there's one thing I'm confusing.

In 'compression/model.py', there is a function, feature_probs_based_sigma.

    def feature_probs_based_sigma(feature, sigma):
        mu = torch.zeros_like(sigma)
        sigma = sigma.clamp(1e-10, 1e10)
        gaussian = torch.distributions.laplace.Laplace(mu, sigma)
        probs = gaussian.cdf(feature + 0.5) - gaussian.cdf(feature - 0.5)
        total_bits = torch.sum(torch.clamp(-1.0 * torch.log(probs + 1e-10) / math.log(2.0), 0, 50))
        return total_bits, probs

In there, the 'gaussian' is defined by 'torch.distributions.laplace.Laplace'. I don't know why you didn't define 'gaussian' by 'torch.distribution.normal.Normal' As far as I understand, in the author's tensorflow implementation he used normal distribution.

If you don't mind, could you answer me?

Thank you again for your contribution.

You Synthesis_prior_net class is called as follow:

    def forward(self, x):
        x = self.relu1(self.deconv1(x))
        x = self.relu2(self.deconv2(x))
        return torch.exp(self.deconv3(x))

Why did you use the exponential function instead of ReLU ? That does not seem to be mentioned in Balle's or your paper.

Hello，I noticed that your code is Balle,2018's work on Pytorch. But it doesn't have the Entropy Model Code and code about using z to AE/AD, In Balle's work, it is done on Tensorflow .Have you finish this part? Thanks for answering.

Do you have any advice regarding MS-SSIM optimization? I get consistently poor results when optimizing directly for the MS-SSIM. Do you have to train longer, mix with other losses, or switch to MS-SSIM later in training in order to get the results described?

Hi jiaheng, I see you use laplace distribution instead of gaussian in the original ICLR2018 paper, is it because laplace is better than gaussian for univariate distribution estimation? so you use GMM in your full model and it's better than laplace?

Hi, thanks so much for sharing your code! I find directly run the train.py may cause too much cpu consumption(~5000% on my machine), so I make some modifications to the original code to reduce cpu burden(->100% on my machine).

1. in train.py, set pin_memory=False

train_dataset = Datasets(train_data_dir, image_size)
train_loader = DataLoader(dataset=train_dataset,
                              batch_size=batch_size,
                              shuffle=True,
                              pin_memory=False, 
                              num_workers=2)

2. in model.py, pre-allocate gpu memory of quant_noise data.

class ImageCompressor(nn.Module):
    def __init__(self, out_channel_N=192, out_channel_M=320):
        super(ImageCompressor, self).__init__()
        self.Encoder = Analysis_net(out_channel_N=out_channel_N, out_channel_M=out_channel_M)
        self.Decoder = Synthesis_net(out_channel_N=out_channel_N, out_channel_M=out_channel_M)
        self.priorEncoder = Analysis_prior_net(out_channel_N=out_channel_N, out_channel_M=out_channel_M)
        self.priorDecoder = Synthesis_prior_net(out_channel_N=out_channel_N, out_channel_M=out_channel_M)
        self.bitEstimator_z = BitEstimator(out_channel_N)
        self.out_channel_N = out_channel_N
        self.out_channel_M = out_channel_M
      
        self.quant_noise_feature = torch.zeros(4, self.out_channel_M, 256 // 16, 256 // 16).cuda()
        self.quant_noise_z = torch.zeros(4, self.out_channel_N, 256 // 64, 256 // 64).cuda()
        self.quant_noise_feature = torch.nn.init.uniform_(torch.zeros_like(self.quant_noise_feature), -0.5, 0.5)
        self.quant_noise_z = torch.nn.init.uniform_(torch.zeros_like(self.quant_noise_z), -0.5, 0.5)

    def forward(self, input_image):            
        feature = self.Encoder(input_image)
        batch_size = feature.size()[0]

        z = self.priorEncoder(feature)

Where can we set bit-per-pixel? For example in your article ( fig 9 ), how did you set different bit-per-pixel? Should we train the model for different bit-per-pixels?

Hi, thanks a lot for your reimplementation of "Variational image compression with a scale hyperprior" , it really helps. I wonder if you could provide the pretrained model as well, because I really need it to start my training. Thanks again.

Hi, thanks for your great work. I tried to reproduce the paper's result. I used your code directly (The "train_lambda": 8192,"out_channel_N": 192,"out_channel_M": 320) , and trained the model with the provided dataset from Flicker. However, the val result on Kodak seems bad. The mse loss is about1.2e-4, the estimated bpp is around 1.8 ( the model has already been trained for at least 100 epoch). I do not know what is wrong with my model? Is there any thing that I missed such as doing some data augment??

Hi jiaheng:

Thanks a lot for the repo, it really helps a lot.

One thing I am interested about is to generate the real bit sequence for the latent features and latent variable z. I read the codes in this repo and it seems you only use the bitestimator.py to estimate the entropy but didn't generate the bit sequence. If people consider the transmission of images, then the bit sequence should be important... Did you implement that before or are there any hints for that? Thanks a lot in advance!

Best, Chenghong