Repository: huggingface/diffusers

Issue:
[Dreambooth] Training script mean loss calculation question
Hello,

I was playing with Dreambooth training and noticed a thing [here](https://github.com/huggingface/diffusers/blob/0025626cd9feca260e05f295f272af4f7b9ce44a/examples/dreambooth/train_dreambooth.py#L606) 

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](https://github.com/huggingface/diffusers/blob/0025626cd9feca260e05f295f272af4f7b9ce44a/examples/dreambooth/train_dreambooth.py#L609)



Discussion:
cc @patil-suraj - could you take a look please? :-) 
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?
One can tweak `prior_loss_weight` to change the weights in the sum of  the loss from instance samples and class samples
cc @williamberman @pcuenca 
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 

```python
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.

Resolving patch:
diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py
--- a/examples/dreambooth/train_dreambooth.py
+++ b/examples/dreambooth/train_dreambooth.py
@@ -732,7 +732,7 @@ def main(args):
                     target, target_prior = torch.chunk(target, 2, dim=0)
 
                     # Compute instance loss
-                    loss = F.mse_loss(model_pred.float(), target.float(), reduction="none").mean([1, 2, 3]).mean()
+                    loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
 
                     # Compute prior loss
                     prior_loss = F.mse_loss(model_pred_prior.float(), target_prior.float(), reduction="mean")
diff --git a/examples/textual_inversion/textual_inversion.py b/examples/textual_inversion/textual_inversion.py
--- a/examples/textual_inversion/textual_inversion.py
+++ b/examples/textual_inversion/textual_inversion.py
@@ -634,7 +634,8 @@ def main():
                 else:
                     raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
 
-                loss = F.mse_loss(model_pred.float(), target.float(), reduction="none").mean([1, 2, 3]).mean()
+                loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+
                 accelerator.backward(loss)
 
                 optimizer.step()


Test patch: