StyleGAN2学習済みモデルを使って任意の画像を生成させる

memo.sugyan.com

の続き。

StyleGAN2 は "mapping network" と "synthesis network" の2つのネットワークで構築されていて、画像の生成を行う synthesis network への入力 dlatents_in を変化させていくことで様々な変化を出せる、というものだった。

前回は mapping network からの出力値を使って 「学習によって上手く生成できるようになった画像」のための dlatents_in の値の間を遷移させるといったことをしていたけど、実際には synthesis network には十分に様々な画像を生成できる能力が獲得されているはず、らしい。 具体的には、アイドルの顔画像だけで学習したモデルでも アイドルの顔以外の画像も生成できるかもしれない、ということ。

以下の論文で様々な実験・検証が行われている。

任意画像を生成するための latent space の学習

要するに synthesis network への入力 dlatents_in(14, 512) の shape を持つ変数 (14 は 256x256サイズの場合の数値) とみなし、それを使って生成される画像が目標画像に近くなるように学習させていけば良い、ということ。

せっかくなので TensorFlow 2.x で動くように書いてみた。

Snapshot から SavedModel への変換

StyleGAN2 の公式実装は TensorFlow 1.x でしか動かない。ので 以前の記事 に書いたように snapshot の .pkl から Generatorの部分だけ取り出して SavedModel の形式に変換して保存する。

output_names = [t.name for t in Gs.output_templates]

with tf.Graph().as_default() as graph:
    outputs = [graph.get_tensor_by_name(name) for name in output_names]
    images = tf.transpose(outputs[0], [0, 2, 3, 1])
    images = tf.saturate_cast((images + 1.0) * 127.5, tf.uint8)
    # save as SavedModel
    builder = tf.compat.v1.saved_model.Builder(save_dir)
    signature_def_map = {
        'synthesis': tf.compat.v1.saved_model.build_signature_def(
            {'dlatents': tf.saved_model.utils.build_tensor_info(outputs[1])},
            {'images': tf.saved_model.utils.build_tensor_info(images),
             'outputs': tf.saved_model.utils.build_tensor_info(tf.transpose(outputs[0], [0, 2, 3, 1]))})
    }
    builder.add_meta_graph_and_variables(
        sess,
        [tf.saved_model.tag_constants.SERVING],
        signature_def_map)
    builder.save()

最終的な画像としての出力は [0, 255]tf.uint8 の値に変換したものだけど、ここではその前段階での synthesis network の出力を NHWC に変換しただけのものを使う。 これは (?, 256, 256, 3)tf.float32 tensor で、 [-1.0, 1.0] の範囲の値を持つとみなして処理される。

Keras layers の構築

TensorFlow 2.x では主に Keras API を使って model の構築 & 学習 をしていくことになる。 tf.keras.layers.Layer を継承した独自の layer を定義していく。

class LatentSpace(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__(input_shape=())
        self.v = self.add_weight(
            shape=(1, 14, 512),
            dtype=tf.float32)

    def call(self, inputs):
        return tf.identity(self.v)


class Synthesis(tf.keras.layers.Layer):
    def __init__(self, model_path):
        super().__init__()
        model = tf.saved_model.load(model_path)
        self.synthesis = model.signatures['synthesis']

    def call(self, inputs):
        return self.synthesis(dlatents=inputs)['outputs']

まずは (1, 14, 512) の変数だけを持つ layer。 入力されてくる inputs は無視して、持っている変数をそのまま出力する。この add_weight で登録された変数たちが、training すべきパラメータとなる。

次に その変数の値を入力として受けて 生成を行う layer。 これは先述した SavedModelload して結果を返してやるだけで良い。

これで Model の構築ができる。

model = tf.keras.Sequential([
    LatentSpace(),
    Synthesis(model_path),
])
model.summary()
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #
=================================================================
latent_space (LatentSpace)   (1, 14, 512)              7168
_________________________________________________________________
synthesis (Synthesis)        (1, 256, 256, 3)          0
=================================================================
Total params: 7,168
Trainable params: 7,168
Non-trainable params: 0
_________________________________________________________________

Target image と Dataset

生成目標となる画像を用意する。普通に読み込んで decode すると [0, 255] 範囲の tf.uint8 tensor になってしまうので、synthesis network の出力に合わせて [-1.0, 1.0] の範囲になるよう調整する。

それを使って 学習時の入力データを作成する。といっても Model への入力は不要なので適当に 0 とかを返しておく。 Target data となる y だけ常に同じ値を返し続ければ良い。

with open(target_image, 'rb') as fp:
    y = tf.image.decode_jpeg(fp.read())
y = tf.expand_dims(tf.cast(y, tf.float32) / 127.5 - 1.0, axis=0)

dataset = tf.data.Dataset.from_tensors((0, y))

Loss class

あとは最小化すべき loss の定義。 論文によると 生成画像と目標画像の間の pixel-wise MSE と、 VGG16 を使った perceptual loss を組み合わせて使うようだ。 要するに pixel 間の差分だけでなく 特徴も似たようなものになるのが良い、ということのようで。

tf.keras.losses.Loss を継承した独自の loss を定義する。

class EmbeddingLoss(tf.keras.losses.Loss):
    def __init__(self, image):
        super().__init__()
        self.vgg16 = tf.keras.applications.VGG16(include_top=False)
        self.target_layers = {'block1_conv1', 'block1_conv2', 'block3_conv2', 'block4_conv2'}
        self.outputs = []
        out = image
        for layer in self.vgg16.layers:
            out = layer(out)
            if layer.name in self.target_layers:
                self.outputs.append(out)

    def call(self, y_true, y_pred):
        out = y_pred
        outputs = []
        for layer in self.vgg16.layers:
            out = layer(out)
            if layer.name in self.target_layers:
                outputs.append(out)
        n = tf.cast(tf.math.reduce_prod(y_pred.shape), tf.float32)
        losses = tf.math.reduce_sum(tf.math.squared_difference(y_true, y_pred)) / n
        for i, out in enumerate(outputs):
            n = tf.cast(tf.math.reduce_prod(out.shape), tf.float32)
            losses += tf.math.reduce_sum(tf.math.squared_difference(self.outputs[i], out)) / n
        return losses

VGG16 の model は tf.keras.applicationsimagenet で学習したものがあるようなので それをそのまま使う。論文によると conv1_1, conv1_2, conv3_2, conv4_2 の4つの layer の出力を使ってそれぞれ差分を足し合わせて loss の値にしている、とのこと。 目標画像については値が変化しないので、この中間層の特徴量も変化しない。ので最初に計算して保持しておく。 call() 時には y_pred で model からの出力値が渡されてくるので、その都度 VGG16 に通して各層の出力を取得する。 それぞれ目標値との tf.math.squared_differencetf.math.reduce_sum して、それぞれの scale で割ってやる。 最終的な和が、最小化すべき loss の値になる。

学習

ここまで出来たらあとは compile して fit させるだけ。 前述の EmbeddingLossloss に指定して、 Adam optimizer で最適化していく。 実験してみた感じではこの optimizer のパラメータによって学習の結果も大きく変わってくるようで、このへんの最適な値を見つけるのはとても難しそうだった。 ここでは論文記載と同じ learning_rate=0.01, epsilon=1e-08 を使用する。

fit では 前述の datasetbatch_size: 1 で繰り返す。適当な stepsで 1 epoch の区切りにして、その epoch 終了時の変数での生成結果を画像として出力するようにする。

class GenerateCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs):
        v = self.model.layers[0].variables[0].numpy()
        images = self.model.layers[1](v)
        images = tf.saturate_cast((images + 1.0) * 127.5, tf.uint8)
        with open(f'epoch{epoch:03d}.png', 'wb') as fp:
            data = tf.image.encode_png(tf.squeeze(images, axis=0)).numpy()
            fp.write(data)


model.compile(
    optimizer=tf.keras.optimizers.Adam(
        learning_rate=0.01,
        epsilon=1e-08),
    loss=EmbeddingLoss(y))
model.fit(
    dataset.repeat().batch(1),
    steps_per_epoch=50,
    epochs=100,
    callbacks=[GenerateCallback()])

さすがにこれは CPU環境ではそれなりに時間がかかって厳しい。 Google Colaboratory の GPU Runtime だと数分で 5,000 stepくらいは完了するようだ。

学習結果

自力で収集した アイドルの顔画像 7,500 枚で ある程度学習したモデルを使用。

まずは 実際にこのモデルによって生成された画像。

f:id:sugyan:20200216220601p:plain

これは既に自らが生成した実績がある画像なので、再現できないとおかしいくらいのものではある。

f:id:sugyan:20200216220059g:plain

意外と「完全に一致」というところまではいかない…。けど まぁ早い段階からほぼ再現できているようには見える。 ここがより安定して近くなるかどうかは optimizer のパラメータ次第という感じではあった。

次に、この生成モデルへの学習にはまったく使っていない 女優さんの顔画像とかだとどうなるだろうか。

f:id:sugyan:20200216220015g:plain

ちょっとハッキリしないけど、一応それなりに髪や顔のパーツまで生成できているようだ。

では もはや日本人でも若い女性でもない人物の顔画像を目標にした場合は…?

f:id:sugyan:20200216220128g:plain

思ったよりイケる! これはこれで予想外。

まったく学習データに使っていない画像も生成できるようになる、というのは面白いな〜。 今回使ったモデルはかなり学習データのドメインが限定的だし 学習も完了ってほど十分に出来ていないのだけど、それでもこれだけ生成できる、ということが分かった。

もっと学習が進んだものや データを増やして学習させたモデルを使った場合にはまた違う結果になるかもしれない。

おまけ: morphing

こうして任意の画像を生成するための synthesis network への入力が得られた、ということは 前回の記事 で書いたように、2つの入力があった場合にその間の値を使うことで morphingが出来る…はず。

と思ってやってみたが

f:id:sugyan:20200216215837g:plain

f:id:sugyan:20200216215904g:plain

f:id:sugyan:20200216215127g:plain

と 中間の表現は気持ち悪いばかりのものになってしまった。

これくらい離れた空間同士だと単純な線形の推移では自然な変化を出せないようだ。このあたりも学習の進行度合いによって違ったりするかもしれないけど…。

Repository