はじめに
本記事は、LOCAL学生部アドベントカレンダー2019の12日目の記事になります。
https://adventar.org/calendars/4020
アドベントカレンダーのネタ探しに時間がかかりすぎてなかなか決まらなかったので、既出ではありますが、Pix2Pixという生成モデルを使用したモノクロ画像の着色にチャレンジしたお話をします。
(Speech Synthesisやりたかった.....)
Pix2Pixとは
arXiv：https://arxiv.org/abs/1611.07004
GANの手法を利用し、様々なドメイン画像間の変換を統一的なネットワーク、誤差関数で行えるようにしたもの。
論文で紹介されているのは、
- マスク画像から元画像
- モノクロ画像からカラー画像
- 航空写真から地図
- 輪郭から色ぬり画像
もちろんこれらを逆に変換するよう訓練することも可能です。
ネットワークアーキテクチャ
Generator
Generatorのネットワークは、以下のようになっています。
U-Netのモデルを修正したオートエンコーダになっています。
Conv -> Batchnorm -> Leaky ReLU
と画像を畳み込んで行き、デコーダーでは
Transposed Conv -> Batchnorm -> Dropout -> ReLU
と言った風に画像を拡大するように畳み込んで行きます。
入力画像とGeneratorの生成画像との絶対値誤差(MAE)と、Discriminator(後述)の出力による交差エントロピー(Sigmoid cross entropy)、この二つを足し合わせてGeneratorの損失関数とします。
Discriminator
Discriminatorの方のネットワークは以下の通りです。
Discriminatorは、Generatorが生成した画像と教師データとして与えられている画像を識別するためのネットワークです。生成画像と訓練画像を入力として、本物か偽物(生成されたもの)かを見極めます。
訓練フェーズでは、GeneratorとDiscriminatorとが互いに競い合うようにして精度を高めていきます。
実装
Tensorflow公式のチュートリアルを参考に実装しました。2.0の波に乗れてなかった人なので新しくなった関数群に戸惑いながら作りました。
今回チャレンジするタスクは、入力としてモノクロ画像を与え、カラー画像を出力させるというタスクです。データセットは、Googleが公開しているOpen Images dataset V5を使いました。
全体の実装はGithubにあります。
https://github.com/kodamanbou/Pix2Pix
Input pipeline
画像データのロードを行い、リサイズ、ランダムクロップ、グレースケール化します。
def load(image_file):
image = tf.io.read_file(image_file)
image = tf.image.decode_jpeg(image)
input_image = tf.image.rgb_to_grayscale(image)
input_image = tf.cast(input_image, tf.float32)
real_image = tf.cast(image, tf.float32)
print(input_image.shape)
print(real_image.shape)
return input_image, real_image
def resize(input_image, real_image, height, width):
input_image = tf.image.grayscale_to_rgb(input_image)
input_image = tf.image.resize(input_image, [height, width],
method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
real_image = tf.image.resize(real_image, [height, width],
method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
return input_image, real_image
def random_crop(input_image, real_image):
stacked_image = tf.stack([input_image, real_image], axis=0)
cropped_image = tf.image.random_crop(
stacked_image, size=[2, hp.image_size, hp.image_size, 3])
return cropped_image[0], cropped_image[1]
def normalize(input_image, real_image):
input_image = (input_image / 127.5) - 1
real_image = (real_image / 127.5) - 1
return input_image, real_image
@tf.function
def random_jitter(input_image, real_image):
# resizing to 286 x 286 x 3
input_image, real_image = resize(input_image, real_image, 286, 286)
# randomly cropping to 256 x 256 x 3
input_image, real_image = random_crop(input_image, real_image)
if tf.random.uniform(()) > 0.5:
# random mirroring
input_image = tf.image.flip_left_right(input_image)
real_image = tf.image.flip_left_right(real_image)
input_image = tf.image.rgb_to_grayscale(input_image)
return input_image, real_image
def load_image_train(image_file):
input_image, real_image = load(image_file)
input_image, real_image = random_jitter(input_image, real_image)
input_image, real_image = normalize(input_image, real_image)
return input_image, real_image
def load_image_test(image_file):
input_image, real_image = load(image_file)
input_image, real_image = resize(input_image, real_image,
hp.image_size, hp.image_size)
input_image = tf.image.rgb_to_grayscale(input_image)
input_image, real_image = normalize(input_image, real_image)
return input_image, real_image
Generator
Generatorは、先ほど説明したU-Net likeなネットワークになります。
def downsample(filters, size, apply_batchnorm=True):
initializer = tf.random_normal_initializer(0., 0.02)
result = tf.keras.Sequential()
result.add(
tf.keras.layers.Conv2D(filters, size, strides=2, padding='same',
kernel_initializer=initializer, use_bias=False))
if apply_batchnorm:
result.add(tf.keras.layers.BatchNormalization())
result.add(tf.keras.layers.LeakyReLU())
return result
def upsample(filters, size, apply_dropout=False):
initializer = tf.random_normal_initializer(0., 0.02)
result = tf.keras.Sequential()
result.add(
tf.keras.layers.Conv2DTranspose(filters, size, strides=2,
padding='same',
kernel_initializer=initializer,
use_bias=False))
result.add(tf.keras.layers.BatchNormalization())
if apply_dropout:
result.add(tf.keras.layers.Dropout(0.5))
result.add(tf.keras.layers.ReLU())
return result
def Generator():
inputs = tf.keras.layers.Input(shape=[256, 256, 1])
down_stack = [
downsample(64, 4, apply_batchnorm=False), # (bs, 128, 128, 64)
downsample(128, 4), # (bs, 64, 64, 128)
downsample(256, 4), # (bs, 32, 32, 256)
downsample(512, 4), # (bs, 16, 16, 512)
downsample(512, 4), # (bs, 8, 8, 512)
downsample(512, 4), # (bs, 4, 4, 512)
downsample(512, 4), # (bs, 2, 2, 512)
downsample(512, 4), # (bs, 1, 1, 512)
]
up_stack = [
upsample(512, 4, apply_dropout=True), # (bs, 2, 2, 1024)
upsample(512, 4, apply_dropout=True), # (bs, 4, 4, 1024)
upsample(512, 4, apply_dropout=True), # (bs, 8, 8, 1024)
upsample(512, 4), # (bs, 16, 16, 1024)
upsample(256, 4), # (bs, 32, 32, 512)
upsample(128, 4), # (bs, 64, 64, 256)
upsample(64, 4), # (bs, 128, 128, 128)
]
initializer = tf.random_normal_initializer(0., 0.02)
last = tf.keras.layers.Conv2DTranspose(3, 4,
strides=2,
padding='same',
kernel_initializer=initializer,
activation='tanh') # (bs, 256, 256, 3)
x = inputs
skips = []
for down in down_stack:
x = down(x)
skips.append(x)
skips = reversed(skips[:-1])
for up, skip in zip(up_stack, skips):
x = up(x)
x = tf.keras.layers.Concatenate()([x, skip])
x = last(x)
return tf.keras.Model(inputs=inputs, outputs=x)
def generator_loss(disc_gen_out, gen_out, target):
gan_loss = loss_object(tf.ones_like(disc_gen_out), disc_gen_out)
l1_loss = tf.reduce_mean(tf.abs(gen_out - target))
total_gen_loss = gan_loss + (hp.LAMBDA * l1_loss)
return total_gen_loss, gan_loss, l1_loss
Discriminator
識別側のネットワークです。
def Discriminator():
initializer = tf.random_normal_initializer(0., 0.02)
inp = tf.keras.layers.Input(shape=[256, 256, 3], name='input_image')
tar = tf.keras.layers.Input(shape=[256, 256, 3], name='target_image')
x = tf.keras.layers.concatenate([inp, tar]) # [N, 256, 256, 4]
down1 = downsample(64, 4, False)(x) # [N, 128, 128, 64]
down2 = downsample(128, 4)(down1) # [N, 64, 64, 128]
down3 = downsample(256, 4)(down2) # [N, 32, 32, 256]
zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3) # [N, 34, 34, 256]
conv = tf.keras.layers.Conv2D(512, 4, strides=1,
kernel_initializer=initializer,
use_bias=False)(zero_pad1) # [N, 31, 31, 512]
batchnorm1 = tf.keras.layers.BatchNormalization()(conv)
leaky_relu = tf.keras.layers.LeakyReLU()(batchnorm1)
zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu) # [N, 33, 33, 512]
last = tf.keras.layers.Conv2D(1, 4, strides=1,
kernel_initializer=initializer)(zero_pad2) # [N, 30, 30, 1]
return tf.keras.Model(inputs=[inp, tar], outputs=last)
def discriminator_loss(disc_real_out, disc_gen_out):
real_loss = loss_object(tf.ones_like(disc_real_out), disc_real_out)
gen_loss = loss_object(tf.zeros_like(disc_gen_out), disc_gen_out)
total_disc_loss = real_loss + gen_loss
return total_disc_loss
Training
定義したネットワークを呼び出し、実際にトレーニング、モデルの保存などを行います。
@tf.function
def train_step(input_image, target, epoch):
with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
gen_output = generator(input_image, training=True)
disc_real_output = discriminator([input_image, target], training=True)
disc_generated_output = discriminator([input_image, gen_output], training=True)
gen_total_loss, gen_gan_loss, gen_l1_loss = generator_loss(disc_generated_output, gen_output, target)
disc_loss = discriminator_loss(disc_real_output, disc_generated_output)
generator_gradients = gen_tape.gradient(gen_total_loss,
generator.trainable_variables)
discriminator_gradients = disc_tape.gradient(disc_loss,
discriminator.trainable_variables)
generator_optimizer.apply_gradients(zip(generator_gradients,
generator.trainable_variables))
discriminator_optimizer.apply_gradients(zip(discriminator_gradients,
discriminator.trainable_variables))
with summary_writer.as_default():
tf.summary.scalar('gen_total_loss', gen_total_loss, step=epoch)
tf.summary.scalar('gen_gan_loss', gen_gan_loss, step=epoch)
tf.summary.scalar('gen_l1_loss', gen_l1_loss, step=epoch)
tf.summary.scalar('disc_loss', disc_loss, step=epoch)
def fit(train_ds, epochs, test_ds):
for epoch in range(epochs):
start = time.time()
for example_input, example_target in test_ds.take(1):
generate_images(generator, example_input, example_target)
print("Epoch: ", epoch)
# Train
for n, (input_image, target) in train_ds.enumerate():
print('.', end='')
if (n + 1) % 100 == 0:
print()
train_step(input_image, target, epoch)
print()
# saving (checkpoint) the model every 20 epochs
if (epoch + 1) % 20 == 0:
checkpoint.save(file_prefix=checkpoint_prefix)
print('Time taken for epoch {} is {} sec\n'.format(epoch + 1,
time.time() - start))
checkpoint.save(file_prefix=checkpoint_prefix)
結果
Epoch 0
Epoch 1
Epoch 15
追加で載せていきます...
まとめ
ほぼチュートリアル通りの実装ですが、他のタスクにも簡単に応用できそうです。
参考サイト
https://www.tensorflow.org/tutorials/generative/pix2pix
http://yusuke-ujitoko.hatenablog.com/entry/2017/06/25/165008