TensorFlow2.0 + 無料のColab TPUでDCGANを実装した

TensorFlow2.0とGoogle Colaboratoryの無料TPUを使って、DCGANを実装しました。

## 訓練経過の様子
<blockquote class="twitter-tweet"><p lang="und" dir="ltr"><a href="https://t.co/2qXowYMk6t">pic.twitter.com/2qXowYMk6t</a></p>&mdash; しこあん@『モザイク除去本』好評通販中 (@koshian2) <a href="https://twitter.com/koshian2/status/1186164682320568320?ref_src=twsrc%5Etfw">October 21, 2019</a></blockquote> <script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script> 

## 何をやったか
* Google ColabのTPU＋TF2.0でCelebA（約20万枚）をDCGANで生成
* TF1.X系のTPUでは、同時に実行可能なグラフは1個の制約があったため、GANの訓練が容易ではなかった（[こちらの記事にある](https://qiita.com/mgmk2/items/4847e01501a3d56a3094)通り、不可能であったわけではない。しかし、低レベルAPIが必須で決して容易ではなかった）。**TF2.X系のTPUでは、もっと容易にGANを実装できた**。
* DCGANの論文通りのモデル（パラメーター数：G=12.7M, D=11.0M）で。**64x64の画像20万枚を、1エポックを40秒程度で訓練可能**。100エポック回しても1時間程度。
* **大きなバッチサイズ（BS=1024）で訓練できた**。BigGANの論文にもあるように、バッチサイズを大きくすることは生成画像の質・安定性の向上ともに重要。
* 1時間程度で以下のような顔画像が生成できた

![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/232088/f618a42a-ca9e-b022-5d50-5d993bdadf69.png)

## Colab Notebook
こちらから遊べます
https://colab.research.google.com/drive/1rSm01ZiAFgxsWOcQ48cdtgLiVnwwZNfx

## TF2.0のTPUの使い方について
基本的な使い方は、MNISTサンプルの記事を書いたのでこちらを参照ください。

[TensorFlow2.0 with KerasでいろいろなMNIST（TPU対応）](https://qiita.com/koshian2/items/5029028fe6d01b243f44)

## 訓練ループの中身
ポイントをかいつまんで解説していきます。

GANの訓練ループをどう書いているかについて。1バッチ単位での訓練は次のようにします。

```python
@distrtibuted(Reduction.SUM, Reduction.SUM, Reduction.CONCAT)
def train_on_batch(real_img1, real_img2):
    # concat x1, x2
    real_img = tf.concat([real_img1, real_img2], axis=0)
    # generate fake images
    with tf.GradientTape() as d_tape, tf.GradientTape() as g_tape:
        z = tf.random.normal(shape=(real_img.shape[0], 1, 1, 100))
        fake_img = model_G(z)
    # train discriminator
    with d_tape:
        fake_out = model_D(fake_img)
        real_out = model_D(real_img)
        d_loss = (loss_func(tf.zeros(shape=(z.shape[0], 1)), fake_out)
                    + loss_func(tf.ones(shape=(z.shape[0], 1)), real_out)) / 2.0
        d_loss = tf.reduce_sum(d_loss) * (1.0 / batch_size)
    gradients = d_tape.gradient(d_loss, model_D.trainable_weights)
    param_D.apply_gradients(zip(gradients, model_D.trainable_weights))
    # train generator
    with g_tape:
        fake_out = model_D(fake_img)
        g_loss = loss_func(tf.ones(shape=(z.shape[0], 1)), fake_out)
        g_loss = tf.reduce_sum(g_loss) * (1.0 / batch_size)
    gradients = g_tape.gradient(g_loss, model_G.trainable_weights)
    param_G.apply_gradients(zip(gradients, model_G.trainable_weights))
    return d_loss, g_loss, fake_img 
```


### distributedデコレーター
これはTensorFlowの組み込みではなく、自分で実装したデコレーターです。

TPUの訓練では（複数GPUと同様）、複数のTPUデバイスにデータをMapしたあと、個々の計算結果（損失値や生成画像）をReduceする必要があります。実装する際、特に意識しないといけないのがReduceで、Reduce用の関数はTFの組み込みでもいくつか用意されています。

しかし、組み込み関数だけでは単体の```train_on_batch```関数を定義したあと、distributed対応の関数を別に書かないといけず、デザイン的に少し野暮ったいのです。そこで、分散訓練の対応＋Reduceをいい感じにやってくれるデコレータを自分で実装しました。詳しくはこちら。

[TensorFlow2.0でDistributed Trainingをいい感じにやるためのデコレーターを作った](https://blog.shikoan.com/distributed-train-decorator-in-tf20/)

デコレーターの引数は各返り値に対するReduceの方法です。```d_loss```と```g_loss```はSUM、```fake_img```はaxis=0でCONCATしています（Concatは組み込みのreduce関数で非対応なので、デコレーターを実装する際は少し工夫がいる）。

### 2つのGradient Tape
GとDという2つのモデルを訓練（偏微分を計算）しないといけないので、Gradient Tapeは2個用意します。これは```tape.gradient()```で偏微分を計算してしまうと、それまでの計算グラフがリセットされてしまうからです。これはPyTorchでも同じです。

TF2.0では、Gradient Tape以下の計算は自動微分可能です。最初の```fake_img```の生成では2個のTapeがあるので、どちらのTapeでもGは微分可能ということになります（もう少し賢い書き方あるかもしれません）。

そして、DのTapeとGのTapeで敵対的学習するようなロスを定義し、学習ステップを進めています。

ここでは生成画像```fake_img```をDとGの間で使いまわししています。**DとGの順番に注意**しましょう。Dを訓練した後はGの係数は変わりませんが、Gを訓練した後は当然Gの係数が変わります。したがって、Gの訓練→Dの訓練での生成画像の使いまわしはできません。**D→Gなら使いまわしできます**。

ちなみにGradient Tapeを2個使って、[2階微分の計算なんかもできたりします](https://blog.shikoan.com/gradienttape-multi/)。WGAN-GPをやりたいときに便利ではないでしょうか。

## データの読み込みについて
GANの訓練よりも**実はここが一番のポイント**だったりします。いくつかトラップがあります。

### TPU+tf.dataではローカルファイルから読み込めない（らしい）
例えば、[tf.data.Dataset.list_files](https://www.tensorflow.org/api_docs/python/tf/data/Dataset#list_files)を使って、「ファイルパス→tf.data内で画像を読み込んでテンソルを返す」という処理は、CPUやGPUでもできてもTPUでは現状はできないようです。

これは[TPUのトラブルシューティング](https://cloud.google.com/tpu/docs/troubleshooting)にも書いてあります（同じのエラーが出ます）。

>ローカル ファイルシステムを使用できない
エラー メッセージ
InvalidArgumentError: Unimplemented: File system scheme '[local]' not implemented

>詳細
すべての入力ファイルとモデル ディレクトリは Cloud Storage バケットパス（gs://bucket-name/...）を使用する必要があり、このバケットは TPU サーバーからアクセス可能である必要があります。すべてのデータ処理とモデル チェックポインティングは、ローカルマシンではなく TPU サーバー上で実行されることに注意してください

これを読む限り、Cloud Storageでないと無理みたいですね。一応、TensorFlowのソースを読んでいくと、[StreamingFilesDataset](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/tpu/datasets.py#L50)というのもあり、コメントを読んでいる限りなんかローカルファイルでも使えそうな気がします。残念ながらドキュメントがなく、どう使うのかはよくわかりませんでした。

そこで今回は、**一度全ての画像を解像度が64x64のNumpy配列に格納し、それを```from_tensor_slices```でtf.dataに読み込ませることで解決**を図りました。

### TensorFlowのオブジェクトの2GBの壁
Numpy配列化すれば全て解決というわけではありません。CelebAをNumpy配列化すると、uint8型でも**202599枚 × 64px × 64px × 3ch = 2.31GB**と大容量になってしまいます。

TensorFlowのオブジェクトには1個あたり2GBの制約があります。```from_tensor_slices```では、内部的に一度Numpy配列を定数のテンソルに置き換えている（と思われる）ので、2GBをオーバーするとエラーが出ます。[公式ドキュメント](https://www.tensorflow.org/api_docs/python/tf/data/Dataset#from_tensor_slices)にも以下のような記載があります。

> Note that if tensors contains a NumPy array, and eager execution is not enabled, the values will be embedded in the graph as one or more tf.constant operations. For large datasets (> 1 GB), this can waste memory and run into byte limits of graph serialization. 

そこで、**今回はデータを半分に分割して、あたかも2つの画像テンソルが流れてくるようなデータセットとみなす**ようにしました。全体のテンソルは2.31GBなので、半分に分割すれば2GBの制約はクリアできます。前半をX1、後半をX2とします。擬似コードですが、

```python
for X1, X2 in dataset:
    # TPUデバイス内で
    X = tf.concat([X1, X2], axis=0)
```

とすれば、データを半分に分割しても、個々のTPUデバイス側で結合することができます（だいぶ頭おかしい解決方法）。これでちゃんと訓練できました。動いてしまえば正義ですね。

ちなみにNumpy配列化したときに、そこまで容量が大きくなければ（2GBを超えなければ）このような心配をする必要はありません。

## 訓練経過
### 1エポック
![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/232088/825fb8db-1c76-31a8-9974-2c6a99ec19d6.png)

### 25エポック
そこそこ形にはなってきました。
![epoch_0024.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/232088/6d8a6b14-acbc-db99-1bfb-71b27c092e0e.png)

### 50エポック
相当綺麗です。バッチサイズが1024と大きいおかげで、勾配の信頼性が高く、DCGANでも安定しやすいのでしょう。ここまで30分程度です。
![epoch_0050.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/232088/005236ec-b8d8-3a2e-6f51-1707511c18b4.png)

### 100エポック
1時間ちょいでこのようになりました。ただのDCGANでここまでいけるのはすごい。
![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/232088/f618a42a-ca9e-b022-5d50-5d993bdadf69.png)

## まとめ
この記事では、TensorFlow2.0＋Colab TPUを使って、CelebA20万枚をDCGANで1時間程度で訓練する方法を紹介しました。

今までは、GANをColab上で訓練することは厳しかったが実情でした。なぜなら、TPUでは訓練コードを書くことが容易ではなかったですし、GPUでは実行時間12時間の制約にかかりやすく、小さい解像度でしか訓練できなかったからです。GANでは自前のGPUが推奨で、遊ぶためのハードルが高いのが現実問題としてありました。

しかし、TensorFlow2.0を使うと、TPUでもGANが容易な形で実装可能（データの読み込みなど多少面倒なところはありますが）となったため、**GANのハードルがぐっと下がった**といえるでしょう。

さらに、TPUの計算はfloat32でも非常に高速で、同じ内容をGPUで行う場合、1ポック40秒程度で回すのは相当なGPU数がいると思われます（まずバッチサイズ1024で訓練するのが相当大変）。それだけの計算資源を無料で得られるわけですから、やはりTPUはすごい。

記事の冒頭にNotebookを公開しているので、ぜひ遊んでみてください。