search

yuval-alaluf / restyle-encoder

Official Implementation for "ReStyle: A Residual-Based StyleGAN Encoder via Iterative Refinement" (ICCV 2021) https://arxiv.org/abs/2104.02699
https://yuval-alaluf.github.io/restyle-encoder/
MIT License
1.03k stars 154 forks source link

The model collapse when use moco loss #50

Closed tommy-qichang closed 2 years ago

tommy-qichang commented 2 years ago

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)

image

tommy-qichang commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

OK. Thanks.

tommy-qichang commented 2 years ago

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 commented 2 years ago

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

tommy-qichang commented 2 years ago

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

yuval-alaluf commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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