StyleGANをGoogle Colaboratoryで実装したい

Question

前提・実現したいこと

Google ColaboratoryでStyleGANを実装しようとしています。(初学者です)
generatorをロードするのurlを変更することで他の画像も生成できると記載があった為、ImageNetデータセットで学習する予定です。（imagenetを用いたstyleganがstylegan+clipなどしかなかったため）

generatorのロードの部分で
ModuleNotFoundError: No module named 'torch_utils.persistence'
というエラーがでています

主に以下の記事を参考にして進めていましたが、わかりませんでした。
・https://teratail.com/questions/295390
・https://qiita.com/pacifinapacific/items/1d6cca0ff4060e12d336
・http://cedro3.com/ai/stylegan/
・https://qiita.com/Phoeboooo/items/12d21916de56d125f0be

発生している問題・エラーメッセージ

from google.colab import drive
drive.mount('/content/drive')
%cd /content/drive/My Drive


#git clone でStyleGANのコードを使えるようにする
!git clone https://github.com/NVlabs/stylegan.git

#ディレクトリ移動
%cd stylegan
!pip uninstall tensorflow -y
!pip install tensorflow-gpu==1.15.0
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()


generatorをロードする

import os
import pickle
import numpy as np
import PIL.Image
import dnnlib
import dnnlib.tflib as tflib
import config

def main():
  # Initialize TensorFlow.
  tflib.init_tf()

  # Load pre-trained network.
  url = 'https://s3.eu-central-1.amazonaws.com/avg-projects/stylegan_xl/models/imagenet512.pkl' # karras2019stylegan-ffhq-1024x1024.pkl
  with dnnlib.util.open_url(url, cache_dir=config.cache_dir) as f:
    _G, _D, Gs = pickle.load(f)
      # _G = Instantaneous snapshot of the generator. Mainly useful for resuming a previous training run.
      # _D = Instantaneous snapshot of the discriminator. Mainly useful for resuming a previous training run.
      # Gs = Long-term average of the generator. Yields higher-quality results than the instantaneous snapshot.

  # Print network details.
  Gs.print_layers()
 
  # Pick latent vector.
  rnd = np.random.RandomState(10)  # seed = 10
  latents0 = rnd.randn(1, Gs.input_shape[1])
  latents1 = rnd.randn(1, Gs.input_shape[1])
  latents2 = rnd.randn(1, Gs.input_shape[1])
  latents3 = rnd.randn(1, Gs.input_shape[1])
  latents4 = rnd.randn(1, Gs.input_shape[1])
  latents5 = rnd.randn(1, Gs.input_shape[1])
  latents6 = rnd.randn(1, Gs.input_shape[1])
    
  num_split = 39  # 2つのベクトルを39分割
  for i in range(40):
      latents = latents6+(latents0-latents6)*i/num_split
      # Generate image.
      fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
      images = Gs.run(latents, None, truncation_psi=0.7, randomize_noise=True, output_transform=fmt)
 
      # Save image.
      os.makedirs(config.result_dir, exist_ok=True)
      png_filename = os.path.join(config.result_dir, 'photo'+'{0:04d}'.format(i)+'.png')
      PIL.Image.fromarray(images[0], 'RGB').save(png_filename)
 
if __name__ == "__main__":
  main()

該当のソースコード

ModuleNotFoundError                       Traceback (most recent call last)
<ipython-input-8-a1bc236cff97> in <module>
     45 
     46 if __name__ == "__main__":
---> 47   main()

<ipython-input-8-a1bc236cff97> in main()
     14   url = 'https://s3.eu-central-1.amazonaws.com/avg-projects/stylegan_xl/models/imagenet512.pkl' # karras2019stylegan-ffhq-1024x1024.pkl
     15   with dnnlib.util.open_url(url, cache_dir=config.cache_dir) as f:
---> 16     _G, _D, Gs = pickle.load(f)
     17       # _G = Instantaneous snapshot of the generator. Mainly useful for resuming a previous training run.
     18       # _D = Instantaneous snapshot of the discriminator. Mainly useful for resuming a previous training run.

ModuleNotFoundError: No module named 'torch_utils.persistence'

補足情報（FW/ツールのバージョンなど）

ImageNetの.plkファイルが見つけれず、下記サイトのモデル選択のurlをダウンロードしてGoogle Driveにアップロードして用いています。
エラーは
"https://s3.eu-central-1.amazonaws.com/avg-projects/stylegan_xl/models/imagenet512.pkl"
が原因だと思うのですが公式サイトなどにImageNetの.pklファイルがありましたら教えてもらえると幸いです。

・https://colab.research.google.com/github/kaz12tech/ai_demos/blob/master/StyleGAN_XL_demo.ipynb#scrollTo=sd3w1cEaIxPN

Comments

[jbpb0]

質問の題名の「解決済」とは、どういう意味でしょうか？

[riku_university]

タイトルの確認不足でした！
ご指摘ありがとうございます。

[jbpb0]

> エラーは
"https://s3.eu-central-1.amazonaws.com/avg-projects/stylegan_xl/models/imagenet512.pkl"
が原因だと思う

上記は「StyleGAN-XL」
https://github.com/autonomousvision/stylegan_xl
用のもののようですが、質問者さんが動かそうとしてる「StyleGAN」
https://github.com/NVlabs/stylegan
でも使えるのでしょうか？

https://ninhydrin.hatenablog.com/entry/2022/02/15/091413
に
「提案するStyleGAN-XLはこれを解決してImageNetのような分散の大きなデータの生成を可能にした。」
と書かれてるので、

> ImageNetデータセットで学習する予定です。

をやるなら、「StyleGAN」ではなくて「StyleGAN-XL」を動かす必要があるのではないですかね

[jbpb0]

https://github.com/justinpinkney/awesome-pretrained-stylegan2#Imagenet
に、「StyleGAN2」
https://github.com/NVlabs/stylegan2
用のImagenet pre-trained modelsがありました

[riku_university]

やはりその部分が原因でしたか・・・StyleGAN-XL用のものでも動くかなと思いましたが無理そうですね。

ImageNeを用いると StyleGAN の性能が落ちるというのは理解していて、どれだけ落ちるのか気になったため StyleGAN に Imageanet を用いて実装しようとしていました。

StyleGAN2 のほうを試しに Google Colab で実装してみようと思います。URLありがとうございます。

Answer 1

エラーは
"https://s3.eu-central-1.amazonaws.com/avg-projects/stylegan_xl/models/imagenet512.pkl"
が原因だと思う

上記は
StyleGAN-XL
用のもののようなので、質問者さんが動かそうとしてる
StyleGAN
では使えないのではないですかね