train with stylegan2_pytorch model Gs.pth throw KeyError: 'g_ema' error

[t109598032]

Hi, thank you for providing the code. I firstly train my own dataset on stylegan2_pytorch(made by adriansahlman), then it will generate Gs.pth as the stylegan2 generator weights file, then I put Gs.pth as the stylegan_weights of pSp, then it throw the error output like this:

/opt/conda/lib/python3.6/site-packages/setuptools/distutils_patch.py:26: UserWarning: Distutils was imported before Setuptools. This usage is discouraged and may exhibit undesirable behaviors or errors. Please use Setuptools' objects directly or at least import Setuptools first.
  "Distutils was imported before Setuptools. This usage is discouraged "
{'batch_size': 8,
 'board_interval': 50,
 'checkpoint_path': None,
 'dataset_type': '3dp_encode',
 'encoder_type': 'GradualStyleEncoder',
 'exp_dir': 'exp_dir2',
 'id_lambda': 0,
 'image_interval': 100,
 'input_nc': 3,
 'l2_lambda': 1.0,
 'l2_lambda_crop': 0,
 'label_nc': 0,
 'learn_in_w': False,
 'learning_rate': 0.0001,
 'lpips_lambda': 0.8,
 'lpips_lambda_crop': 0,
 'max_steps': 100000,
 'moco_lambda': 0,
 'optim_name': 'ranger',
 'output_size': 1024,
 'resize_factors': None,
 'save_interval': 5000,
 'start_from_latent_avg': True,
 'stylegan_weights': 'pretrained_models/Gs.pth',
 'test_batch_size': 8,
 'test_workers': 4,
 'train_decoder': True,
 'use_wandb': False,
 'val_interval': 2500,
 'w_norm_lambda': 0,
 'workers': 4}
Loading encoders weights from irse50!
Loading decoder weights from pretrained!
Traceback (most recent call last):
  File "scripts/train.py", line 33, in <module>
    main()
  File "scripts/train.py", line 28, in main
    coach = Coach(opts)
  File "./training/coach.py", line 36, in _init_
    self.net = pSp(self.opts).to(self.device)
  File "./models/psp.py", line 34, in _init_
    self.load_weights()
  File "./models/psp.py", line 63, in load_weights
    self.decoder.load_state_dict(ckpt['g_ema'], strict=False)
KeyError: 'g_ema'

Does it anything wrong? thank you very much!

[yuval-alaluf]

Please see the details provided here on correctly preparing your generator: https://github.com/yuval-alaluf/restyle-encoder#preparing-your-generator (The same details also apply here to pSp)

[YKJamesMoriarty]

@t109598032 Have you solved this problem, I also encountered the same problem, I read the author's answer but did not help