1

I trained a StyleGan2-ADA model on a custom dataset which generated a .pkl file. I'm now trying to load the .pkl file so that I can convert it to a .pt file, but when I load the .pkl file using:

pickle.load(f)

I'm getting a ModuleNotFoundError: No module named 'torch_utils.persistence'

I've installed torch_utils and other dependencies, but for loading the file I'm not sure how to fix this issue. If anyone has had this issue in loading a .pkl file any help would be greatly appreciated!!

Same issue on Github here, but no clear solution.

Have tried installing torch_utils multiple times, but error still persists

Naveen M
  • 11
  • 1

1 Answers1

1

Include dnnlib and torch_utils packages to your project, you can find it at https://github.com/NVlabs/stylegan3.

import pickle
import torch
import dnnlib
import torch_utils

with open('your.pkl', 'rb') as f:
    G = pickle.load(f)['G_ema'].cpu()  # torch.nn.Module
z = torch.randn([1, G.z_dim]).cpu()  # latent codes
c = None  # class labels (not used in this example)
img = G(z, c)
user941581
  • 389
  • 2
  • 4