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

memo.sugyan.com

の続き。

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

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

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

Image2StyleGAN: How to Embed Images Into the StyleGAN Latent Space?

任意画像を生成するための 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。これは先述した SavedModel を load して結果を返してやるだけで良い。

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

学習

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

fit では前述の dataset を batch_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()])