SinGANを一部実装してみた

##概要
　[この論文](https://arxiv.org/abs/1905.01164)で提案されているSinGANを実装してみました。
Qiitaにも[[論文解説】SinGAN: Learning a Generative Model from a Single Natural Image](https://qiita.com/takoroy/items/27f918a2fe54954b29d6)や[SinGANの論文を読んだらテラすごかった](https://qiita.com/yoyoyo_/items/81f0b4ca899152ac8806)といった解説記事を書かれている方もいます。

　上記の記事でも解説されていますがSinGANは超解像・アニメーション・Paint2Imageなど様々なタスクを同じモデル構造で学習できる手法です。今回はその中で超解像を実装しました。記事のタイトルに一部と付けているのはそのためです。

コードは以下にあります。
https://github.com/Sakai0127/SinGAN_tf_impl

学習にはgooglecolabを使用しました。ノーノブックへのリンクはこちら。
https://colab.research.google.com/drive/12HNbuFSfpKb2D0bZt080KJpgnJvc4ZWo

##コード解説
　使用するフレームワークはtensorflowです。モデルの解説というよりtensorflowとkerasで実装する際にハマった部分の解説が主になります。

　その前に学習結果を。左が学習に使った低解像の画像で、右が生成した高解像の画像です。ここでは約4倍のサイズに拡大しています。

|LR|HR|
|:-:|:-:|
|![33039_LR.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/194245/469545ba-88b5-cf7e-c419-96dc3e0d1c41.png)|![479x319.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/194245/c448cdc7-596d-ab40-7550-cbc7a31a0e09.jpeg)|

 論文や解説記事を読めばわかりますが、SinGANの構造はかなりシンプルです。Convolution+BatchNormalization+LeakyReluを繰り返していくだけです。ただし、Generatorの最終層だけはTanhを使い入力と足し合わせます。コードは以下のような感じで。

```python
def conv_block(inputs, name='conv_block', layer_idx=0, out_channels=32, activation=tf.keras.layers.LeakyReLU(0.2), bn=True):
    x = tf.keras.layers.Conv2D(out_channels, 3, padding='SAME', name='%s_conv_%d'%(name, layer_idx))(inputs)
    if bn:
        x = tf.keras.layers.BatchNormalization(name='%s_BN_%d'%(name, layer_idx))(x)
    x = activation(x)
    return x

def G_block(inputs, noise, name='G_block', hidden_maps=32, num_layers=5, out_channel=3):
    with tf.name_scope(name):
        x = inputs + noise
        x = conv_block(x, name='conv_block_0', layer_idx=0, out_channels=hidden_maps)
        for i in range(1, num_layers-1, 1):
            x = conv_block(x, name='conv_block_%d'%i, layer_idx=i, out_channels=hidden_maps)
        x = conv_block(x, name='conv_block_%d'%(num_layers-1), layer_idx=num_layers-1, out_channels=out_channel, activation=tf.keras.activations.tanh, bn=False)
    return x + inputs

def D_block(inputs, name='D_block', hidden_maps=32, num_layers=5):
    with tf.name_scope(name):
        x = conv_block(inputs, name='conv_block_0', layer_idx=0, out_channels=hidden_maps)
        for i in range(1, num_layers-1, 1):
            x = conv_block(x, name='conv_block_%d'%i, layer_idx=i, out_channels=hidden_maps)
        x = conv_block(x, name='conv_block_%d'%(num_layers-1), layer_idx=num_layers-1, out_channels=1, activation=tf.keras.activations.linear, bn=False)
    return x
```

　このGとDのブロックを解像度ごとに積み重ねていくだけです。nameをかなり細かく指定していますが、そのあたりはもっとスマートにしたいです。各Gブロックでは入出力は共に同じサイズの画像となっています。低解像から順に学習していくという性質上、次の層への入力はアップサンプルする必要があります。それはtf.image.resizeを使って行っています。また、論文に従ってadversarial loss関数はWGAN-gpを使っています。
　ここまでは特に問題ないと思いますが、トレーニングループを行う箇所で詰まりました。1回のパラメータ更新は以下のコードで行います。

```python
@tf.function
def train_step(real, g_input, noise_w, G, D, opt_G, opt_D, g_times, d_times, alpha):
    for _ in range(d_times):
        with tf.GradientTape() as d_tape:
            noise = tf.random.normal(tf.shape(real))*noise_w
            d_loss = util.calc_d_loss(real, g_input, noise, D, G, alpha)
        d_grad = d_tape.gradient(d_loss, D.trainable_weights)
        opt_D.apply_gradients(zip(d_grad, D.trainable_weights))
    for _ in range(g_times):
        with tf.GradientTape() as g_tape:
            noise = tf.random.normal(tf.shape(real))*noise_w
            g_loss = util.calc_g_loss(real, g_input, noise, D, G, alpha)
        g_grad = g_tape.gradient(g_loss, G.trainable_weights)
        opt_G.apply_gradients(zip(g_grad, G.trainable_weights))
    return g_loss, d_loss
```

　ここで問題になったのはapply_gradientの箇所です。最初のブロックでは問題ありませんが、次のブロックを追加して学習を開始すると以下のようなエラーが発生しました。同じような問題は[ここ](https://github.com/tensorflow/tensorflow/issues/27120)で報告されています。

```
ValueError: tf.function-decorated function tried to create variables on non-first call.
```


###tf.keras.optimizers.Adam
　上記のエラーメッセージを読むと「初回呼び出し以外で@tf.function内で変数を作成しようとした」ということのようです。ですが、新しい層を追加する際に新規でOptimizerインスタンスを生成するようにしていたのでエラーの意味が分かりませんでした。現在のtensorflow(2.0.0)では複数のoptimizerを使い分けようとすると上記エラーが発生してしまうようです。今回は以下のように対処しました。@tf.functionの外で変数を作成するわけです。

```python
def init_opt(opt, model):
    opt.iterations
    opt._create_hypers()
    opt._create_slots(model.trainable_weights)
```

　kerasのOptimizerでは初回のapply_gradientの呼び出しでパラメータ更新のための関数が作成されます。今回はAdam Optimizerを使用しましたが、他の手法でも同様に使えると思います。Adamの場合はパラメータ数は2n+1個になります。2nの部分は各変数に対してm, v 2種類のパラメータを用いて最適化を行うためです。これらを作成しているのがopt._create_slots(model.trainable_weights)の部分です。また、kerasのOptimizerではiteration回数を保持する変数も作成されます。+1の部分がそれでopt.iterationsで作成しています。opt._create_hypers()はハイパーパラメータを作成しています。learning_rateはすべてのOptimizerで共通ですがAdamの場合はbeta1, beta2, epsilonも含まれます。これらの変数を事前に作成することで問題なくトレーニングを行えました。

　今回はSuper Resolutionを実装しましたが、SinGANは他にも色々できるのでそちらも試してみたです。