How can I convert Imagenet checkpoint to a .pkl file to run stylegan2-ada-pytorch?

[bekir23]

Hi. the paper that you wrote is quite useful for gan training. But I try to use Imagenet checkpoints that you share to run stylegan2-ada-pytorch and then this repo ask me .pkl file to start training. I exported [StyleGAN-ADA-256]. It include folder of 'data', 'data.pkl' and 'version' files. This data.pkl file's size is too little. I think it is pointless to start gan training because similar model files has 100mb or up size. I think that I have to convert 'data' folder to 'data.pkl'. How can I do this?

[anvoynov]

Hi! Do you use the file from the link from the readme?
https://www.dropbox.com/s/7gll7weysn1ull7/imagenet-256-state.pt
it's the standard pytorch *.pt file and you can use it with the official StyleGAN2-ADA code: https://github.com/NVlabs/stylegan2-ada-pytorch

[bekir23]

These .pt files do not be used in https://github.com/NVlabs/stylegan2-ada-pytorch . It asks .pkl file. I have to convert pt file to pkl fıle. I searched a litle bit on it and encountered some codes that it show me a way to solve it (https://pytorch.org/tutorials/beginner/saving_loading_models.html)
According to above link. I decide to create model structure to load .pt file.

I added some codes in training_loop.py in the stylegan2-ada repo.
training_loop.py load structure of model in below lines while training:

G = dnnlib.util.construct_class_by_name(**G_kwargs, **common_kwargs).train().requires_grad_(False).to(device) # subclass of torch.nn.Module
D = dnnlib.util.construct_class_by_name(**D_kwargs, **common_kwargs).train().requires_grad_(False).to(device) # subclass of torch.nn.Module
G_ema = copy.deepcopy(G).eval()

Thus, we reach 3 model structure and the we should load our weight to these strucrutes. First we load pt file and matching appropiate weights to related model structure:

pt_file = torch.load('/stylegan2-ada-pytorch/pretrained_models/imagenet-256-state.pt')
G.load_state_dict(pt_file['G'])
D.load_state_dict(pt_file['D'])
G_ema.load_state_dict(pt_file['G_ema'])

Then we have to save this model as pickle file. Pickle file consist of dict objects,create dict then load your model into it:

resume_data={}
for name, module in [('G', G), ('D', D), ('G_ema', G_ema)]:
    resume_data[name]=module

then save it as pickle:

dest='/stylegan2-ada-pytorch/pretrained_models/imagenet-256-state.pkl'
with open(dest, 'wb') as f:
    pickle.dump(resume_data, f)

Then you can use this pkl file to train your gan model in styklegan2 ada repo.

If do you know any better solution to use these weights in stylegan2 ada repo. Please let me know.
Thank you

[naeun0620]

Can you upload **G_kwargs, **D_kwargs, and **common_kwargs too? @bekir23