openclimatefix / skillful_nowcasting

Implementation of DeepMind's Deep Generative Model of Radar (DGMR) https://arxiv.org/abs/2104.00954
MIT License
211 stars 59 forks source link

Wrong dimension to split when training Discriminator #27

Closed PingChiaHuang closed 2 years ago

PingChiaHuang commented 2 years ago

It seems that there are some bugs when training the discriminator.

The score of real inputs and generated inputs are splitted from "concatenated_outputs" along dimension 1. https://github.com/openclimatefix/skillful_nowcasting/blob/35e25ce553e49937f26e88c0f2bb2b36adfb77bc/dgmr/dgmr.py#L155 However, the discriminator concat the spatial loss and temporal loss along dimension 1, forming "concatenated_outputs". https://github.com/openclimatefix/skillful_nowcasting/blob/35e25ce553e49937f26e88c0f2bb2b36adfb77bc/dgmr/discriminators.py#L37

Therefore, "score_real" and "score_generated" actually represent the spatial loss and temporal loss instead of what you desire. I will appreciate it if you can tell me your opinion :)

jaggbow commented 2 years ago

I noticed the same thing and was about to open an issue too until I saw yours. Indeed, the splitting should be done batch-wise instead:

Old line: score_real, score_generated = torch.split(concatenated_outputs, 1, dim=1)

Suggested change:

B = concatenated_outputs.shape[0]
score_real, score_generated = torch.split(concatenated_outputs, B, dim=0)
score_real_spatial, score_real_temporal = torch.split(score_real, 1, dim=1)
score_generated_spatial, score_generated_temporal = torch.split(score_generated, 1, dim=1)
discriminator_loss = loss_hinge_disc(score_generated_spatial, score_real_spatial)+loss_hinge_disc(score_generated_temporal, score_real_temporal)
jacobbieker commented 2 years ago

Ah okay, thanks for pointing this out, I'll push the fix soon!