はじめに

Google Colabで2022年8月でTF1のサポートが終了し、TF1で作成されたモデルの実行ができなくなりました。これにより当法人でGANの勉強用のコンテンツとして使っていた StyleGAN2もGoogle Colab上で動かなくなってしまいました（✳︎）。
本記事では、StyleGAN2を新しいTF2上のGoogle Colabで動かした一例を紹介します(2022年10月19日現在での確認であり今後とも本対処で動くとは限りません）。

他のTF1で作成されたモデルについても簡易的に新しいGoogle Colab上で実行する際の参考になればと思います。

記載事項の誤り等がありましたら随時ブラッシュアップしていきますのでご指摘いただければと思います。

✳︎補足：あくまでGoogle Colabで動かなくなっただけであり、個別のGPUサーバ上での環境でTF1の環境では問題なく動きます。Google Colab上では、tensorflow-gpuをinstallすれば少なくても2022年９月中は動かせたのですが、それさえも動かなくなりました（GPUの割り当てが何故かされません）。

なお、StyleGAN2の説明については、下記の記事で詳細に記載されていますので参考になるかと思います。

学習コードの修正

注意事項として、画像の生成時は学習時のコードもロードされるためStyleGAN2のOfficial GitHubにある事前学習済みモデルは使うことができず、学習データを用意して学習を一から実行し直す必要があります（精度が高い画像生成を行うには専用のGPUサーバ等で行う必要があります）。

TF2上で動かす方針としてTF2に完全に移行するのではなく、TF２の環境上でTF1互換で動作させます。
なお、StyleGAN2のモデルは config-f を対象に検証しています。

以下に順を追って説明します。

`training_loop`

import文の修正と、TF2でcontrib.memory_statsは削除されたためメモリ使用レポートの値はゼロにします。レポートのみなので学習には直接影響しないためです。
参考記事：https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md

training/training_loop.py

#import tensorflow as tf
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
・・・
def training_loop( 
    ...
    #peak_gpu_mem_op = tf.contrib.memory_stats.MaxBytesInUse()
    peak_gpu_mem_op = tf.constant(0)
・・・

`tfutil.py`

import文の修正とcontributeの使用をコメントアウトします。

dnnlib/tflib/tfutil.py

#import tensorflow as tf
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
・・・
#import tensorflow.contrib   # requires TensorFlow 1.x!
#tf.contrib = tensorflow.contrib

custom_ops.py

import文の修正と、カスタムオペレーションの nvcc(cudaコンパイル) 時のオプションを修正します。
下記の記事を参照しました。
https://stackoverflow.com/questions/71827660/undefined-symbol-at-runtime

dnnlib/tflib/custom_ops.py


#import tensorflow as tf
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
    ・・・
   #compile_opts += ' --compiler-options \'-fPIC -D_GLIBCXX_USE_CXX11_ABI=0\''
   compile_opts += ' --compiler-options \'-fPIC -D_GLIBCXX_USE_CXX11_ABI=1\''

`optimizer.py`

import文の修正と、使用するオプティマイザをTF1互換のオプティマイザに修正します。
下記の記事を参考にしました。
https://www.tensorflow.org/api_docs/python/tf/compat/v1/train/Optimizer

dnnlib/tflib/optimizer.py

#import tensorflow as tf
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
・・・

    #tf_optimizer: str = "tf.train.AdamOptimizer", # Underlying optimizer class.
    tf_optimizer: str = "tf.compat.v1.train.AdamOptimizer",
・・・

その他

Tensorflowを使うモジュールのimport文の修正をします。

training/dataset.py

#import tensorflow as tf
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
・・・

training/loss.py

#import tensorflow as tf
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
・・・

training/network_stylegan2.py

#import tensorflow as tf
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
・・・

dnnlib/tflib/autosummary.py

#import tensorflow as tf
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
・・・

dnnlib/tflib/network.py

#import tensorflow as tf
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
・・・

これで準備完了です（データセットは事前に作成してください）。

run_trainingを起動しますが、起動時のパラメータで　--metrics=noneを指定する必要があります。
デフォルトでは fid50kが使用されますが、各種metrics用の事前学習済みモデルはTF1で学習されているためロードできません。そのため、metricsは使用せずに学習を行います。なおmetricsは学習時の評価のために使用されているので学習そのものには影響はありません。

起動パラメータの例

python run_training.py --num-gpus=1 --result-dir=results/learn1 --data-dir=datasets --config=config-f --metrics=none --dataset=learndata --mirror-augment=true --total-kimg=学習数

生成コードの修正

基本的に学習コードの修正が行われていれば生成コードでの修正はありません。
生成の起動パラメータで学習時に生成された学習済みモデル（.pkl）を指定します。

python run_generator.py generate-images --network=新たに学習したモデル.pkl  --seeds=6600-6605 --truncation-psi=0.5

おわりに

TF１からTF2の移行について移行ツールがあるもののハードルが高いと思われます。
今後、TF1で作成された各モデルをTF2へ本格移行にチャレンジし得られたノウハウを記事にしていきたいと思います。

StyleGAN2をTF2オンリーになったGoogle Colabで動かしてみた