The model collapse when use moco loss

[tommy-qichang]

Hi Yuval,

When I train my datasets with the configuration you recommended, the model seems to collapse after just a few iterations(As shown below). However, when I set the moco lambda as 0.0001, there is no problem so far. So do you have any clue why that happens? Is there a way to disable the moco loss totally? (when I set the moco loss as 0, it seems you bind the id_logs with the moco_loss calculation)

[tommy-qichang]

I notice in the paths_config.py, you are using

'moco': 'pretrained_models/moco_v2_800ep_pretrain.pth.tar'

However, in the README, google drive link download the pretrained_models/moco_v2_800ep_pretrain.pt
Is there any differences between them?
I download the pth.tar from Facebook's GitHub(rather than use the pt file) and it leads to the error.

[yuval-alaluf]

Using the MoCo model that was uploaded to the Google Drive should be the correct version. It is possible that the config is incorrect and the correct path should be 'moco': 'pretrained_models/moco_v2_800ep_pretrain.pt'.
I will recheck this.

[tommy-qichang]

OK. Thanks.

[tommy-qichang]

Sorry. it seems the pt version of Moco still has the same problem.
Did you have such a problem before? and is there any way to disable moco ? Thanks.

[yuval-alaluf]

Can you please provide the error you're getting with the MoCo model?

[tommy-qichang]

There is no error, but the synthetic image seems not the real MRI image. And the loss didn't decrease.

[yuval-alaluf]

Oh. I don't think your problem is the MoCo loss then. You should be able to disable it completely by setting the lambda value to 0. But I don't think this will necessarily solve your problem.
Do you know if your inversion problem is even solvable using an encoder? For example, have you tried overfitting ReStyle on a small set of images (e.g., less than 5)? Whenever I start working on a new domain, I perform these small sanity checks to make sure that the problem is even solvable.
In addition, have you tried performing the inversion using optimization? If optimization can't solve the inversion accurately, I don't believe ReStyle will be able to solve your task.
Finally, how is your StyleGAN generator? Is it able to produce high-quality scans like your input and target images shown to the left? ReStyle is limited to what StyleGAN can generate.
These are some of the questions I would ask before tackling the losses in the training process.
Hope this helps!

[tommy-qichang]

I've been using pixel2style2pixel for few times with the model, and it works well. The only unsatisfaction is the inversion details are not as good as the optimization method(The optim method can perfectly invert the real image).

[yuval-alaluf]

It is interesting that the results you got with pSp are better than ReStyle. Did you use the same loss weights in both cases?

[tommy-qichang]

I just use l2 and lpips loss. And I've tried using id loss, which shows a comparable result. But when I add moco loss, the synthetic image starts to become worse. I'm not sure why, but I'll keep my eyes on the problem. Thank you.

[yuval-alaluf]

It does make some sense that the moco loss may lead to unsatisfactory results because MoCo was pretrained on ImageNet which is quite different from your data. Have you tried disabling the moco loss by setting the lambda value to 0?

[tommy-qichang]

Yes, I'm trying to disable Moco loss and right now the training goes well.
Thank you so much for your guidance.