StyleGAN2をGoogle Colaboratoryで実装したい

Question

前提・実現したいこと

Google Colaboratory で StyleGAN2 を試しています。(初学者です)
10月のどこかで CUDA10.0 が消えたりと Google Colab に改正があり、その影響で元は動いていたコードが動かなくなってしまいした。

以下のサイトを参考にして進めています。
・https://techpr.info/ml/stylegan2-colab/

発生している問題①

下記のコードを実行するとエラーがでます。

vector_size = Gs.input_shape[1:][0]
seeds = seeder(range(8000, 8020), vector_size)

generate_images(Gs, seeds, truncation_psi=0.5)

エラーメッセージ①

Generating image for seed 0/20 ...
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/client/session.py in _do_call(self, fn, *args)
   1364     try:
-> 1365       return fn(*args)
   1366     except errors.OpError as e:

8 frames
InvalidArgumentError: Cannot assign a device for operation Gs/_Run/Gs/latents_in: {{node Gs/_Run/Gs/latents_in}} was explicitly assigned to /device:GPU:0 but available devices are [ /job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:XLA_CPU:0, /job:localhost/replica:0/task:0/device:XLA_GPU:0 ]. Make sure the device specification refers to a valid device.
	 [[Gs/_Run/Gs/latents_in]]

During handling of the above exception, another exception occurred:

InvalidArgumentError                      Traceback (most recent call last)
/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/client/session.py in _do_call(self, fn, *args)
   1382                     '\nsession_config.graph_options.rewrite_options.'
   1383                     'disable_meta_optimizer = True')
-> 1384       raise type(e)(node_def, op, message)
   1385 
   1386   def _extend_graph(self):

InvalidArgumentError: Cannot assign a device for operation Gs/_Run/Gs/latents_in: node Gs/_Run/Gs/latents_in (defined at /usr/local/lib/python3.7/dist-packages/tensorflow_core/python/framework/ops.py:1748)  was explicitly assigned to /device:GPU:0 but available devices are [ /job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:XLA_CPU:0, /job:localhost/replica:0/task:0/device:XLA_GPU:0 ]. Make sure the device specification refers to a valid device.
	 [[Gs/_Run/Gs/latents_in]]

発生している問題②

生成した画像（image○○○.png）を動画形式に変換するコードでエラーが出ます。

!ffmpeg -r 30 -i image%d.png -vcodec mpeg4 -y test.mp4

エラーメッセージ②

[image2 @ 0x55fd18c8e000] Could find no file with path 'image%d.png' and index in the range 0-4
image%d.png: No such file or directory

該当のソースコード

# install cuda10.0
!apt-get --purge remove cuda nvidia* libnvidia-*
!dpkg -l | grep cuda- | awk '{print $2}' | xargs -n1 dpkg --purge
!apt-get remove cuda-*
!apt autoremove

!wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-repo-ubuntu1804_10.0.130-1_amd64.deb
!sudo dpkg -i cuda-repo-ubuntu1804_10.0.130-1_amd64.deb
!sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
!sudo apt-get update
!wget http://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb
!sudo apt install -y ./nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb
!sudo apt-get update

# Install NVIDIA driver
#!sudo apt-get install --no-install-recommends nvidia-driver-418
!sudo apt-get -y installnvidia-driver-418
# Reboot. Check that GPUs are visible using the command: nvidia-smi

# Install development and runtime libraries (~4GB)
#!sudo apt-get install --no-install-recommends \
!sudo apt-get install -y \
    cuda-10-0 \
    libcudnn7=7.6.2.24-1+cuda10.0  \
    libcudnn7-dev=7.6.2.24-1+cuda10.0


# Install TensorRT. Requires that libcudnn7 is installed above.
# !sudo apt-get install -y --no-install-recommends libnvinfer5=5.1.5-1+cuda10.0 \
!sudo apt-get install -y libnvinfer5=5.1.5-1+cuda10.0 \
    libnvinfer-dev=5.1.5-1+cuda10.0

!apt --fix-broken install


# tensorflow
!pip uninstall tensorflow -y
!pip install tensorflow-gpu==1.15.2
!pip install keras==2.2.4


# Git clone
!git clone https://github.com/NVlabs/stylegan2.git


# 出力先フォルダの作成
!mkdir out


# 潜在変数を作成するための関数を用意
def seeder(seeds, vector_size):
  result = []

  for seed in seeds:
    rnd = np.random.RandomState(seed)
    result.append( rnd.randn(1, vector_size) ) 
  return result


# 必要なパッケージのインストール
import argparse
import numpy as np
import PIL.Image
import dnnlib
import dnnlib.tflib as tflib
import re
import sys
import pretrained_networks


# 画像生成プロセスの関数
def generate_images(Gs, seeds, truncation_psi):
    noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]

    Gs_kwargs = dnnlib.EasyDict()
    Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
    Gs_kwargs.randomize_noise = False
    if truncation_psi is not None:
        Gs_kwargs.truncation_psi = truncation_psi

    for seed_idx, seed in enumerate(seeds):
        print('Generating image for seed %d/%d ...' % (seed_idx, len(seeds)))
        rnd = np.random.RandomState()
        tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width]
        images = Gs.run(seed, None, **Gs_kwargs) # [minibatch, height, width, channel]
        path = f"/content/out/image{seed_idx}.png"
        PIL.Image.fromarray(images[0], 'RGB').save(path)


# 顔画像の学習モデルを読み込み
network_pkl = 'gdrive:networks/stylegan2-ffhq-config-a.pkl'
print('Loading networks from "%s"...' % network_pkl)
_G, _D, Gs = pretrained_networks.load_networks(network_pkl)


# パラメータの指定
vector_size = Gs.input_shape[1:][0]
seeds = seeder(range(8000, 8020), vector_size)

generate_images(Gs, seeds, truncation_psi=0.5)


# 画像を300枚生成
seeds = seeder([6660, 7770], vector_size)

STEPS = 300
diff = seeds[1] - seeds[0]
step = diff / STEPS
current = seeds[0].copy()

seeds2 = []
for i in range(STEPS):
  seeds2.append(current)
  current = current + step

generate_images(Gs, seeds2, truncation_psi=0.5)


# 生成した画像を動画形式に変換
!ffmpeg -r 30 -i image%d.png -vcodec mpeg4 -y test.mp4

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

GPUを認識できていると思うのですが、なぜ使用できないのでしょうか。

[name: "/device:CPU:0"
 device_type: "CPU"
 memory_limit: 268435456
 locality {
 }
 incarnation: 2440703754907394877, name: "/device:XLA_CPU:0"
 device_type: "XLA_CPU"
 memory_limit: 17179869184
 locality {
 }
 incarnation: 9295743096350903082
 physical_device_desc: "device: XLA_CPU device", name: "/device:XLA_GPU:0"
 device_type: "XLA_GPU"
 memory_limit: 17179869184
 locality {
 }
 incarnation: 3392596144835826399
 physical_device_desc: "device: XLA_GPU device", name: "/device:GPU:0"
 device_type: "GPU"
 memory_limit: 14947935847
 locality {
   bus_id: 1
   links {
   }
 }
 incarnation: 428819314228490541
 physical_device_desc: "device: 0, name: Tesla T4, pci bus id: 0000:00:04.0, compute capability: 7.5"]

Comments

[jbpb0]

> # 必要なパッケージのインストール

の直前あたりで

%cd stylegan2

をしておかないと、

> import dnnlib
import dnnlib.tflib as tflib

や

> import pretrained_networks

で、下記のエラーが出るはずです
ModuleNotFoundError: No module named 'dnnlib'
ModuleNotFoundError: No module named 'pretrained_networks'

それなのに、「発生している問題①、②」では、コードのそれよりも後の部分を実行しているので、なぜエラーが出ずにそこまで到達できてるのかが、不思議です

もし上記のインポートのエラーが出ないなら、それでインポートされてる「dnnlib」や「pretrained_networks」は、
https://github.com/NVlabs/stylegan2
のとは別のものですが、そのようなものを別途インストールしてるのでしょうか？

[jbpb0]

「発生している問題①」よりも後の「発生している問題②」を実行する際には、「発生している問題①」のエラーをどのようにして回避してるのでしょうか？

まさか、「発生している問題①」のエラーを放置して、そのままそれよりも後を実行してるのではありませんよね？

[jbpb0]

「エラーメッセージ②」の

> image%d.png: No such file or directory

は、作成した画像ファイルは

> # 出力先フォルダの作成
!mkdir out

にあるのに、そこじゃない場所で

> !ffmpeg -r 30 -i image%d.png -vcodec mpeg4 -y test.mp4

を実行してるためです

上記を実行する直前に「出力先フォルダ」に「%cd」する必要があります

[riku_university]

履歴を遡ったら以前①のエラーが出なかったのは
import sys
sys.path.insert(0, "/content/stylegan2")
が記載されていたからのようでした。

下記を実行するのに必要なかったため消していたようです。

# モデルのダウンロード
!wget https://battle.shawwn.com/sdc/stylegan2-imagenet-512/model.ckpt
# 画像生成の実行
!python /content/stylegan2/run_generator.py generate-images \
--network=model.ckpt-533504.pkl \
--seeds=6600-6625 --truncation-psi=1.0

[riku_university]

「発生している問題②」は10月以前の改正が入る前まではその手前まで動いていたため、そのときに出たエラーを記載しました。

①が解決しても②が出てくるのは分かっていたので一緒に記載させていただきました。

[jbpb0]

インポートよりも前(上)に

> import sys
sys.path.insert(0, "/content/stylegan2")

を実行したら、

%cd stylegan2

は要りません

その場合は、

> # 生成した画像を動画形式に変換

のすぐ上に、

%cd out

または

%cd /content/out

を追加

[jbpb0]

私の回答の通りにやって、実行できませんか？

[riku_university]

BA選択したはずが押せてなかったようです・・・

Answer 1

質問のコードで下記を変更したら、質問に書いてるエラーは出ませんでした

　

# 必要なパッケージのインストール

のすぐ上に下記を追加

%cd stylegan2

　

# 生成した画像を動画形式に変換

のすぐ上に下記を追加

%cd ../out

　
【追記】
一つ目の変更をしないと、

import dnnlib

の実行で、「ModuleNotFoundError: No module named 'dnnlib'」というエラーが出るはずです
「発生している問題①」は、そのエラーを放置して続きを実行したのが原因かもしれません

　
二つ目の変更をしないと、作成した画像ファイルが無いところで

!ffmpeg -r 30 -i image%d.png -vcodec mpeg4 -y test.mp4

を実行するので、「発生している問題②」になります