概要
この論文で提案されているSinGANを実装してみました。
Qiitaにも[論文解説】SinGAN: Learning a Generative Model from a Single Natural ImageやSinGANの論文を読んだらテラすごかったといった解説記事を書かれている方もいます。
上記の記事でも解説されていますがSinGANは超解像・アニメーション・Paint2Imageなど様々なタスクを同じモデル構造で学習できる手法です。今回はその中で超解像を実装しました。記事のタイトルに一部と付けているのはそのためです。
コードは以下にあります。
https://github.com/Sakai0127/SinGAN_tf_impl
学習にはgooglecolabを使用しました。ノーノブックへのリンクはこちら。
https://colab.research.google.com/drive/12HNbuFSfpKb2D0bZt080KJpgnJvc4ZWo
コード解説
使用するフレームワークはtensorflowです。モデルの解説というよりtensorflowとkerasで実装する際にハマった部分の解説が主になります。
その前に学習結果を。左が学習に使った低解像の画像で、右が生成した高解像の画像です。ここでは約4倍のサイズに拡大しています。
| LR | HR |
|---|---|
 |
 |
論文や解説記事を読めばわかりますが、SinGANの構造はかなりシンプルです。Convolution+BatchNormalization+LeakyReluを繰り返していくだけです。ただし、Generatorの最終層だけはTanhを使い入力と足し合わせます。コードは以下のような感じで。
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回のパラメータ更新は以下のコードで行います。
@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の箇所です。最初のブロックでは問題ありませんが、次のブロックを追加して学習を開始すると以下のようなエラーが発生しました。同じような問題はここで報告されています。
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の外で変数を作成するわけです。
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は他にも色々できるのでそちらも試してみたです。