Loss of aux_loss keep rising during finetune

[TNTWEN]

Hi! @fracape

Recently i am reproducing a E2E video compression paper based on compressai. Folllow https://github.com/InterDigitalInc/CompressAI/blob/master/examples/train.py ,during traing i keep the learning rate of aux_optimizer the same. But i found the aux_loss keep rising when i finetune the model（which means reduce other parameters optimizer’s learning rate and don't change the learning rate of aux_optimizer）. Even if I try to reduce the learning rate of aux_optimizer , aux loss will still keep rising. But when I encode it into an actual bitstream, the actual bpp does not differ much from the predicted bpp. Can I ignore the rise of aux_loss? Or how to solve it.

And this is the bpp and psnr difference between compress and forward. The difference of psnr is already talked in #41

[TNTWEN]

I only monitored the value of aux_loss at the beginning of training and finetuning, and did not observe aux_loss in the middle stage of training. And at the beginning of training, aux_loss can converge quickly .

[dq0309]

The same phenomenon in my training, do you solve it?

[TNTWEN]

@dq0309
Under my test, I think it was caused by adding noise during training. For MeanScaleHyperprior ,maybe we could use y_hat=quantize_ste(y - means_hat) + means_hat to replace the operation of adding noise and at the same time ,we could make the output of inference consistent with that of compress/decompress . But for JointAutoregressiveHierarchicalPriors ,maybe there is no way to solve it , just ignore it.

[fracape]

Hi,
In my experiments, the aux loss usually converges like in the screenshot below:

Are you training from scratch or fine-tuning a pre-trained model ?
You seem to compress videos, but mention the exemplary train.py loop for images. Are you using examples/train_video.py, which includes default aux_learning_rate?

[TNTWEN]

@fracape Hi!
I train my model from scratch. I only use the LR strategy in train.py which sets aux's LR 1e-3 ,others' LR 1e-4. In the last few epochs of training，i keep aux's LR 1e-3,others' LR 1e-5 to fine-tune my model.

I only monitored the aux_loss at the beginning and end of training. At the beginning ,aux_loss converges rapidly from 12K to nearly 100.In the last few epochs, I found aux_loss rising slowly.But the increase is not too large，nearly from 30 to 38. In the picture you show, the vertical axis value is very large, so the change of aux in the later stage of training may not be obvious.

I tried to adjust aux's LR,but it didn't work. In the next experiment, maybe I need to monitor the change of the aux_loss in the whole process

[fracape]

Here is when you zoom at the end of the screenshot I shared previously. The models are trained using a learning rate scheduler as explained [here](https://interdigitalinc.github.io/CompressAI/zoo.html). We don't use exactly the exemplary training example provided. However, at each plateau the aux loss does not necessarily decreases as you are experiencing.