oneThousand1000 / HairMapper

(CVPR 2022) HairMapper: Removing Hair from Portraits Using GANs.
279 stars 39 forks source link

How is the file StyleGAN2-ada-Generator.pth generated? #21

Open ASAmbitious opened 5 months ago

ASAmbitious commented 5 months ago

I want to ask a question, how is the file StyleGAN2-ada-Generator.pth generated? I have run the stylegan2-ada-pytorch project, and the generated weight file is pkl. If I use pkl directly, it seems not feasible. How can I generate StyleGAN2-ada-Generator.pth with my own data set?

oneThousand1000 commented 5 months ago

Hi!

Please enter the styleGAN2_ada_model dir and run the following code:

import os
from stylegan2_ada import dnnlib
import numpy as np
import torch
import pickle
from stylegan2_ada import legacy
from stylegan2_ada.training.networks import Generator

if __name__ == "__main__":
    with dnnlib.util.open_url('./pretrain/stylegan2-ffhq-1024x1024.pkl') as f:
         G = legacy.load_network_pkl(f)['G_ema'].to(device)

    checkpoint_path = os.path.join('./pretrain/Stylegan2-Generator.pth')
    torch.save(G.state_dict(), checkpoint_path)

It helps you transform the pkl file into pth file which could be loaded by our code.

ASAmbitious commented 5 months ago

Thank you very much for your answer. I have now encountered a new problem. I used stylegan2-ada-pytorch to train a 10,000-round two-dimensional character model. I used the method you told me to convert pkl into pth. Now I have encountered the model. The file is missing keys. Do you have any solution ideas? I will send you a piece of my training data, and this is my error message: ============= Dataset will be saved to ./training_runs/dataset/D2 ============= Initializing attribute classifier. /home/su2204/code/HairMapper/classifier/src/config.py:8: YAMLLoadWarning: calling yaml.load() without Loader=... is deprecated, as the default Loader is unsafe. Please read https://msg.pyyaml.org/load for full details. self._dict = yaml.load(self._yaml) /home/su2204/miniconda3/envs/HairMapper/lib/python3.8/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and will be removed in 0.15, please use 'weights' instead. warnings.warn( /home/su2204/miniconda3/envs/HairMapper/lib/python3.8/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or None for 'weights' are deprecated since 0.13 and will be removed in 0.15. The current behavior is equivalent to passing weights=ResNeXt50_32X4D_Weights.IMAGENET1K_V1. You can also use weights=ResNeXt50_32X4D_Weights.DEFAULT to get the most up-to-date weights. warnings.warn(msg) Loading classification_model generator... Initializing generator. [2024-04-09 17:18:05,019][WARNING] Load model from /home/su2204/code/HairMapper/ckpts/my_ckpt/Stylegan2-Generator1W.pth [2024-04-09 17:18:05,020][INFO] Loading pytorch model from /home/su2204/code/HairMapper/ckpts/my_ckpt/Stylegan2-Generator1W.pth. Loading pytorch model from /home/su2204/code/HairMapper/ckpts/my_ckpt/Stylegan2-Generator1W.pth. Traceback (most recent call last): File "step1_generate_data.py", line 159, in main() File "step1_generate_data.py", line 62, in main model = StyleGAN2adaGenerator(model_name, logger=None, truncation_psi=args.truncation_psi) File "/home/su2204/code/HairMapper/styleGAN2_ada_model/stylegan2_ada_generator.py", line 46, in init super().init(model_name, logger) File "/home/su2204/code/HairMapper/styleGAN2_ada_model/base_generator.py", line 98, in init self.load() File "/home/su2204/code/HairMapper/styleGAN2_ada_model/stylegan2_ada_generator.py", line 64, in load self.model.load_state_dict(torch.load(self.model_path), strict=True) File "/home/su2204/miniconda3/envs/HairMapper/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1604, in load_state_dict raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format( RuntimeError: Error(s) in loading state_dict for Generator: Missing key(s) in state_dict: "mapping.fc2.weight", "mapping.fc2.bias", "mapping.fc3.weight", "mapping.fc3.bias", "mapping.fc4.weight", "mapping.fc4.bias", "mapping.fc5.weight", "mapping.fc5.bias", "mapping.fc6.weight", "mapping.fc6.bias", "mapping.fc7.weight", "mapping.fc7.bias". 0_3 0_8 My stylegan training data is similar to these

oneThousand1000 commented 5 months ago

Could you please check that if your character model actually contains "mapping.fc2.weight", "mapping.fc2.bias", "mapping.fc3.weight", "mapping.fc3.bias", "mapping.fc4.weight", "mapping.fc4.bias", "mapping.fc5.weight", "mapping.fc5.bias", "mapping.fc6.weight", "mapping.fc6.bias", "mapping.fc7.weight", "mapping.fc7.bias"?

You can check that by adding print(G.state_dict().keys()) to the code I just posted.

ASAmbitious commented 5 months ago

This is indeed the case. This part of my file: 'mapping.w_avg', 'mapping.fc0.weight', 'mapping.fc0.bias', 'mapping.fc1.weight', 'mapping.fc1.bias']), is different from the part of your StyleGAN2-ada-Generator.pth: 'mapping.w_avg', 'mapping.fc0.weight', 'mapping.fc0.bias', 'mapping.fc1.weight', 'mapping.fc1.bias', 'mapping.fc2.weight', 'mapping.fc2.bias', 'mapping.fc3.weight', 'mapping.fc3.bias', 'mapping.fc4.weight', 'mapping.fc4.bias', 'mapping.fc5.weight', 'mapping.fc5.bias', 'mapping.fc6.weight', 'mapping.fc6.bias', 'mapping.fc7.weight', 'mapping.fc7.bias']).I only have 0-1, you have 0-7 . Do I have to use the stylegan in the HairMapper project? Because I saw you mentioned that you modified stylegan. I also used your stylegan, but got the error: Traceback (most recent call last): File "/data1/miniconda3/envs/stylegan-ada/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap fn(i, *args) File "/data1/syp/styleGAN2_ada_model/stylegan2_ada/train.py", line 383, in subprocess_fn training_loop.training_loop(rank=rank, args) File "/data1/syp/styleGAN2_ada_model/stylegan2_ada/training/training_loop.py", line 166, in training_loop img = misc.print_module_summary(G, [z, c]) File "/data1/syp/styleGAN2_ada_model/stylegan2_ada/torch_utils/misc.py", line 212, in print_module_summary outputs = module(inputs) File "/data1/miniconda3/envs/stylegan-ada/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl result = forward_call(args, kwargs) File "/data1/syp/styleGAN2_ada_model/stylegan2_ada/training/networks.py", line 711, in forward img, styleSpace_latent,F_code = self.synthesis(ws, *synthesis_kwargs) File "/data1/miniconda3/envs/stylegan-ada/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl result = forward_call(args, **kwargs) File "/data1/syp/styleGAN2_ada_model/stylegan2_ada/training/networks.py", line 605, in forward assert sum([len(i) for i in StyleSpace_latent])==26 AssertionError Is the length of this StyleSpace_latent due to the image or the processing process?

oneThousand1000 commented 5 months ago

Hi, StyleSpace_latent is not used in this project, don't worry.

It seems that your mapping network only contains two layers... If you stiil want to use it, you could try to modify https://github.com/oneThousand1000/HairMapper/blob/ec4f31afe116878eedc0a761a85dcdbe76436753/styleGAN2_ada_model/stylegan2_ada_generator.py#L61 as:

mapping_kwargs      = {"num_layers":2},

to load your own .pth file.

ASAmbitious commented 5 months ago

Thank you very much for your answer, I have completed the data generation. I have a question, why do I have only two generated mapping_kwargs, but yours has 7? Did you modify the code?

oneThousand1000 commented 5 months ago

I have no idea about it. Maybe it is caused by the config that you used to train stylegan2-ada-pytorch.

As you can see here: https://github.com/NVlabs/stylegan2-ada-pytorch/blob/d72cc7d041b42ec8e806021a205ed9349f87c6a4/training/networks.py#L180 the default num_layers is set as 8 in the official stylegan2-ada-pytorch. I used the same config in my project.

Btw, the mapping network is used to transform the z latent into w latent. If you lowered the depth of the network from 8 to 2, it might have some influence on forming the w latent space.

ASAmbitious commented 5 months ago

OK, thank you very much for your patient answer.

ASAmbitious commented 5 months ago

我对此一无所知。也许是由您用于训练 stylegan2-ada-pytorch 的配置引起的。

正如您在这里看到的: https://github.com/NVlabs/stylegan2-ada-pytorch/blob/d72cc7d041b42ec8e806021a205ed9349f87c6a4/training/networks.py#L180 在官方 stylegan2-ada-pytorch 中默认的 num_layers 设置为 8。我在我的项目中使用了相同的配置。

顺便说一句,映射网络用于将 z 潜在值转换为 w 潜在值。如果将网络深度从 8 降低到 2,可能会对 w 潜在空间的形成产生一些影响。

When I used stylegan2-ada-pytorch to train, num_layers was indeed 8. I didn't make any changes. I don't know why the generated pkl was 2. The same problem occurred when converting pkl to pth. I thought I would try it first and see the effect.