[Dreambooth] Training script mean loss calculation question

[zetyquickly]

Hello,

I was playing with Dreambooth training and noticed a thing here

Why did the author use this way of calculation mean of the loss? Is there any particular technical need to do so comparing to this

[patrickvonplaten]

cc @patil-suraj - could you take a look please? :-)

[krahnikblis]

zomg i had exactly the same question while researching the same block of code - it appears to be copypasta from 2 different authors, or one author who likes to switch it up a bit whilst coding... both methods deliver the same result of calculation, though i'm unsure if one has performance benefits?

demo:
import torch
import torch.nn.functional as F

test1 = torch.randn((1,4,64,64))
test2 = torch.randn((1,4,64,64))
mse_goofy_way = F.mse_loss(test1.float(), test2.float(), reduction="none").mean([1, 2, 3]).mean()
mse_easy_way = F.mse_loss(test1.float(), test2.float(), reduction="mean")

mse_goofy_way == mse_easy_way

tensor(True)

---

i don't know if it's cool to augment OP's question? i'm trying to figure out why the "prior preservation loss" calculation is not done against the "class prior" but just a split of the example error in the train step - split the error, MSE both halves, then add back together with a weight on one... is that right? unsure how splitting error of a given step (which could be either the class or the target) would connect any dots between class and target... the original paper describes the loss conceptually, and provides a bunch of mathematical formula notation, but i'm too stupid to read that stuff so can't validate that the python code matches the mathsy-scripts...
net: how does splitting the error, performing MSE on the splits, summing, contribute to the class-to-target learning?

[zetyquickly]

One can tweak prior_loss_weight to change the weights in the sum of the loss from instance samples and class samples

[patrickvonplaten]

cc @williamberman @pcuenca

[patil-suraj]

Hey @zetyquickly , I added it that way because this is how it was done in the original codebase, cf https://github.com/CompVis/stable-diffusion/blob/main/ldm/models/diffusion/ddpm.py#L307

But as @krahnikblis said

mse_goofy_way = F.mse_loss(test1.float(), test2.float(), reduction="none").mean([1, 2, 3]).mean()
mse_easy_way = F.mse_loss(test1.float(), test2.float(), reduction="mean")

both of these behave the same way, will open a PR to correct this and run some experiments to make sure.

[patil-suraj]

PR is here #1861