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

StyleGAN2学習済みモデルを使ったmorphing、latent spaceの探求

学習データはまだまだ収集途中だし 学習もまだ完了とは言えない状態なのだけど、なんとなくそれっぽい顔画像は生成できるくらいまでは出来てきているので、それを使った実験をしてみている。

主にこの記事を参考にしています。

qiita.com

1. latents_in の線形移動

まず最初に試したのは、通常の generator network を使った場合の latents_inを使ったもの。

generator は、(?, 512) の任意の入力を受けて画像を生成するようになっている。 これが所謂 潜在空間 (latent space) というやつで ここの値をうまく選ぶことで所望の出力を得られるようになったりする、というわけだ。

例えば適当な乱数で (2, 512) を作ってそれを入力すると、2つの異なる画像が出力される。

その2つの (1, 512) の乱数ベクトルの、中間の値を入力として使えば 生成される画像も2つの画像の中間のものになるだろう、という考え方。

import numpy as np
import tensorflow as tf
from PIL import Image

model_path = '...'
model = tf.saved_model.load(model_path)
generator = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

rnd = np.random.RandomState(0)
z = rnd.randn(2, 512)

inputs = []
steps = 10
for i in range(10):
    inputs.append(z[0] + (z[1] - z[0]) * i / steps)

for i, latents in enumerate(inputs):
    images = generator(latents=tf.constant([latents], tf.float32))['images']
    Image.fromarray(images.numpy()[0]).save(f'out_{i:02d}.png')

z[0] から z[1] までを 0.1 刻みで線形に変化させてそれぞれ画像を生成した結果。

f:id:sugyan:20200209192014p:plain

右端と左端は(一応)違う顔だが、その間で徐々に変化していっているのが分かる、と思う。 GIFアニメーションにするとこんな感じ。

f:id:sugyan:20200209192737g:plain

2. mapping networkの出力 dlatents_in を使う

StyleGAN の generator は、"mapping network" と "synthesis network" の2つの network によって作られている。 generator に入力された latents_in(?, 512) の値は mapping network で (?, 14, 512) といった もう少し次元の大きい変数に変換される。(14 というのは 256x256 のときの数値で upsample の layer 数が増えるとまた 1618 に変化したりする、のかな。) で、この Disentangled latents と呼ばれる出力が、後段の synthesis network への入力に使われ、実際の画像の合成が行われる、ということになる。 dlatents_in と呼ばれるこの値が 画像生成のための入力としてはより直接的なものになるのかもしれない。

ザックリした理解では、 synthesis network は実に広範囲な生成の能力を持っているが より所望の(学習データに近い)主力を得るための dlatents_in を作り出すための前段として mapping network が作用している、という感じか。実際には その出力からさらに truncation_psi trick といった より質の良い出力を得ることが出来る値にするための工夫がされているようだ。

ともかく、この 2つの network を分けて考えると、2つの画像間を補完する場合も mapping network への入力 latents_in を線形に変化させるよりも synthesis network への入力 dlatents_in を変化させた方が自然なものになりそう。

generator も2つに分けて各 network の入出力を扱えるようにした。 StyleGAN2 では return_dlatents = True と option を指定することで dlatents_intensor も得ることが出来る。

この dlatents_in をまず先に計算し、先程と同様にその値を 0.1 刻みで線形に変化させるようにする。

import numpy as np
import tensorflow as tf
from PIL import Image

model_path = '...'
model = tf.saved_model.load(model_path)
mapping = model.signatures['mapping']

rnd = np.random.RandomState(0)
z = rnd.randn(2, 512)
z = mapping(latents=tf.constant(z, tf.float32))['dlatents'].numpy()

inputs = []
steps = 10
for i in range(10):
    inputs.append(z[0] + (z[1] - z[0]) * i / steps)

synthesis = model.signatures['synthesis']
for i, latents in enumerate(inputs):
    images = synthesis(dlatents=tf.constant([latents], tf.float32))['images']
    Image.fromarray(images.numpy()[0]).save(f'out_{i:02d}.png')

f:id:sugyan:20200209200947p:plain

先述の latents_in の変化だと途中で顎のあたりに手のようなものが現れたりガチャガチャと忙しい変化になってしまっていたのに対し、今度はシームレスに2画像間を推移するようになった。

f:id:sugyan:20200209201028g:plain

latent space を探る

dlatents_in は本当に広い空間で、mapping network を通さずに本当に random な値を使って synthesis network に入力すると まったく汚い出力になってしまう。

f:id:sugyan:20200209201952p:plain

逆にここを上手く最適化してやることで 学習データにまったく存在していない画像も生成することが出来る可能性があるようだ。これについてはまたこれから実験していきたい。


dlatents_in の値を mapping network の出力に絞ることで それなりに学習データに近い出力が得られるようにはなるが、学習データの偏りや学習不足などのせいか どうしても印象が似たような顔ばかりになりやすい。

とはいえ、これも mapping network の出力として得られる dlatents_in が大きく異なっていればそれなりには異なる出力画像になるのではないか?

1000件ほどの random な latents_in を mapping network に入力し、得られた各 dlatents_in の間の距離を計算してみる。これが最も遠い組み合わせを選んで使ったら、出力画像も大きく異なるものになるだろうか?

rnd = np.random.RandomState(0)
z = rnd.randn(1000, 512)
z = mapping(latents=tf.constant(z, tf.float32))['dlatents'].numpy()

distances = []
for i in range(z.shape[0]):
    for j in range(i + 1, z.shape[0]):
        distances.append([np.linalg.norm(z[i] - z[j]), (i, j)])

_, (i, j) = sorted(distances, reverse=False)[0]
z = [z[i], z[j]]

inputs = []
steps = 10
for i in range(10):
    inputs.append(z[0] + (z[1] - z[0]) * i / steps)

f:id:sugyan:20200209204147p:plain

f:id:sugyan:20200209204203g:plain

背景がゴチャゴチャしたり 画像の質はあまり良くないかもしれないけど、それなりにインパクトのある変化をする morphing が出来たような気がする。

ちなみに最も距離が近い2つを選択すると

f:id:sugyan:20200209204427p:plain

f:id:sugyan:20200209204446g:plain

となり、ほぼ違いが感じられないような変化になったので この感覚はまぁ合ってそうだ。

Repository

StyleGAN2による画像生成をCPU環境/TensorFlow.jsで動かす

memo.sugyan.com

の続き。

ようやくTensorFlow.jsを使ってブラウザ上で動かせるようになったので、そのためにやったことメモ。

(まだまだ画像の質とかパフォーマンスの問題とかは色々ある)

CPU環境で動かす

最終的にはTensorFlow.jsでブラウザ上で動かすことが目標だったので別にCPUで動かせる必要は無かったのだけど、どうもGPU環境でしか動かない特殊なOpなどもあってそれが変換後のモデルで実行時にエラーを引き起こす原因だったりするため、まずはCPU環境でも安定して動くようにするのが確実なようだ。

前回記事にも書いた通り、 StyleGAN2 の学習の過程で出力される .pkl ファイルは、CPU環境では読み込むことも出来ない。

import pickle
from dnnlib import tflib

tflib.init_tf()
with open('network.pkl', 'rb') as fp:
    _, _, Gs = pickle.load(fp, encoding='latin1')

Gs.print_layers()
...

RuntimeError: NVCC returned an error. See below for full command line and output log:

nvcc "/Users/sugyan/.ghq/github.com/NVlabs/stylegan2/dnnlib/tflib/ops/fused_bias_act.cu" ...

/bin/sh: nvcc: command not found

そりゃCUDAが入っていない環境ではこうなる。 ので、無理矢理にコードを書き換えてCUDAを使わない reference implementation を利用するようにする。 特に upfirdn_2d の方は幾つかの場所から呼ばれるので、デフォルト引数をまとめて書き換えるのも良いけど impl_dict を書き換えて reference implementation を強制してしまうのがラク

diff --git a/dnnlib/tflib/ops/fused_bias_act.py b/dnnlib/tflib/ops/fused_bias_act.py
index 52f6bfd..c294277 100755
--- a/dnnlib/tflib/ops/fused_bias_act.py
+++ b/dnnlib/tflib/ops/fused_bias_act.py
@@ -63,7 +63,7 @@ def fused_bias_act(x, b=None, axis=1, act='linear', alpha=None, gain=None, impl=

     impl_dict = {
         'ref':  _fused_bias_act_ref,
-        'cuda': _fused_bias_act_cuda,
+        'cuda': _fused_bias_act_ref,
     }
     return impl_dict[impl](x=x, b=b, axis=axis, act=act, alpha=alpha, gain=gain)

diff --git a/dnnlib/tflib/ops/upfirdn_2d.py b/dnnlib/tflib/ops/upfirdn_2d.py
index fd23777..1df2935 100755
--- a/dnnlib/tflib/ops/upfirdn_2d.py
+++ b/dnnlib/tflib/ops/upfirdn_2d.py
@@ -57,7 +57,7 @@ def upfirdn_2d(x, k, upx=1, upy=1, downx=1, downy=1, padx0=0, padx1=0, pady0=0,

     impl_dict = {
         'ref':  _upfirdn_2d_ref,
-        'cuda': _upfirdn_2d_cuda,
+        'cuda': _upfirdn_2d_ref,
     }
     return impl_dict[impl](x=x, k=k, upx=upx, upy=upy, downx=downx, downy=downy, padx0=padx0, padx1=padx1, pady0=pady0, pady1=pady1)

これで 引数が impl='cuda' で指定されてこようが関係なく ref の方が使われるようになる。

こうすると .pkl を読み込むのは問題なくできるようになる。 ただここから実際に計算を実行しようとすると問題が起きるわけで。

Gs.run()GPU device を必要とするので output tensor を指定して tf.Session.run() してみる。

import pickle
import numpy as np
import tensorflow as tf
from dnnlib import tflib

tflib.init_tf()
with open('network.pkl', 'rb') as fp:
    _, _, Gs = pickle.load(fp, encoding='latin1')

graph = tf.get_default_graph()
inputs = graph.get_tensor_by_name(f'Gs/{Gs.input_names[0]}:0')
outputs = graph.get_tensor_by_name(f'Gs/{Gs.output_names[0]}:0')

rnd = np.random.RandomState(0)
z = rnd.randn(1, *Gs.input_shape[1:])
print(tf.get_default_session().run(outputs, feed_dict={inputs: z}))
2020-02-04 23:26:46.471886: E tensorflow/core/common_runtime/executor.cc:642] Executor failed to create kernel. Invalid argument: Conv2DCustomBackpropInputOp only supports NHWC.
         [[{{node Gs/G_synthesis/8x8/Conv0_up/conv2d_transpose}}]]

...

StyleGAN2のモデルはGPU環境で学習する前提で作られているので、その環境に最適化された処理が幾つかある。そのうちの一つが NCHW のdata formatを使っていること。 この形でmodelが作られていると、CPU環境では NHWC にしか対応していないので計算を実行することが出来ないようだ。

1. Graphを書き換える

ということで困った、という記事を書いたところ、以下のようなフィードバックをいただいた。

なるほどー、構築されたGraphを舐めていって inputsoperation を書き換えることで NCHWNHWC に変換する方法があるのか…!

ということで上記を参考にしながら自分で書いてみた。

tflib.init_tf()
with open('network.pkl', 'rb') as fp:
    _, _, Gs = pickle.load(fp, encoding='latin1')

graph = tf.get_default_graph()

target_ops = []
for op in graph.get_operations():
    if not op.name.startswith('Gs/'):
        continue
    if 'data_format' in op.node_def.attr and op.node_def.attr['data_format'].s == b'NCHW':
        target_ops.append(op)

まずは Gs/ 以下の、 'NCHW' という data_format attribute を持つ operation をすべて抽出する。これらが書き換える対象となる。

for target_op in target_ops:
    print(f'op: {target_op.name} ({target_op.type})')
op: Gs/G_synthesis/4x4/Conv/Conv2D (Conv2D)
op: Gs/G_synthesis/4x4/ToRGB/Conv2D (Conv2D)
op: Gs/G_synthesis/8x8/Conv0_up/conv2d_transpose (Conv2DBackpropInput)
op: Gs/G_synthesis/8x8/Conv0_up/Conv2D (Conv2D)
op: Gs/G_synthesis/8x8/Conv1/Conv2D (Conv2D)
op: Gs/G_synthesis/8x8/Upsample/Conv2D (Conv2D)
op: Gs/G_synthesis/8x8/ToRGB/Conv2D (Conv2D)
op: Gs/G_synthesis/16x16/Conv0_up/conv2d_transpose (Conv2DBackpropInput)
op: Gs/G_synthesis/16x16/Conv0_up/Conv2D (Conv2D)

...

それらの operation に対し、まずは inputs のshapeを変換していく。

for target_op in target_ops:
    # Input tensors
    if target_op.type == 'Conv2D':
        inputs = [
            tf.transpose(
                target_op.inputs[0],
                [0, 2, 3, 1],
                name=f'{target_op.name}_input_transpose'),
            target_op.inputs[1]
        ]
    elif target_op.type == 'Conv2DBackpropInput':
        inputs = [
            tf.gather(
                target_op.inputs[0],
                [0, 2, 3, 1],
                name=f'{target_op.name}_output_shape_transpose'),
            target_op.inputs[1],
            tf.transpose(
                target_op.inputs[2],
                [0, 2, 3, 1],
                name=f'{target_op.name}_value_transpose')
        ]

実際に見てみると分かるが、こうして 'NCHW' な op を抽出してみると すべて typeConv2DConv2DBackpropInput のどちらかしかない。 少なくとも StyleGAN2 では、この2つのtypeに対してそれぞれ対応するだけで良い、ということになる。

Conv2D の場合、 inputs0 番目に入力のTensorが入ってくる。これが例えば (?, 1, 11, 11) だったりして (N, C, H, W) に対応する。 ので、 tf.transpose[0, 2, 3, 1] を指定することでこの入力Tensor(N, H, W, C) に変換することが出来る。 1 番目の入力は filter の値のようで、これはdata formatに依存しないのでこのまま使えば良い。

Conv2DBackpropInput の場合はもう少し厄介。どうやら入力は output_shape, filter, value という順番で来るらしい。 1 番目は Conv2D の場合と同様そのまま使って 2 番目が入力Tensorなので やはり同様に [0, 2, 3, 1] で transpose してやる。 そして 0 番目が その出力結果のshapeをどうするか指定するという役割のようで、そのshapeを示す (4,)Tensorが入ってくる。 これは例えば [1, 512, 9, 9] といった形の値で やはり NCHW ならその形に指定するわけだけど、ここではこの op を NHWC に変えてやりたいので、この output_shape も書き換えてやらないと その後の出力の型が合わなくなってしまう。 tf.gather を使ってこの output_shape の中身の順番を入れ替える。

これが出来たら次は attributes。

    # Attributes
    attrs = {}
    for k, v in target_op.node_def.attr.items():
        if k == 'data_format':
            continue
        if target_op.type == 'Conv2DBackpropInput' and k == 'strides':
            strides = v.list.i
            attrs[k] = tf.AttrValue(list=tf.AttrValue.ListValue(i=[
                strides[0],
                strides[2],
                strides[3],
                strides[1],
            ]))
        else:
            attrs[k] = v

各 operation は入力値とは別に? attributes というものを持っているようで、これも書き換えてやる必要がある。

data_formatNCHW である、という情報はここに含まれているので、変換する際にはこれを捨ててしまうことで defaultの NHWC にすることが出来る。

もう一つ Conv2DBackpropInput の場合に変更する必要があるのが strides の値で、これも NCHW のときと NHWC のときで扱いが変わるものらしい。

この strides と前述の output_shape については upfirdn_2d.upsample_conv_2d() に分岐が書かれている。

upfirdn_2d.py 抜粋:

    # Determine data dimensions.
    if data_format == 'NCHW':
        stride = [1, 1, factor, factor]
        output_shape = [_shape(x, 0), outC, (_shape(x, 2) - 1) * factor + convH, (_shape(x, 3) - 1) * factor + convW]
        num_groups = _shape(x, 1) // inC
    else:
        stride = [1, factor, factor, 1]
        output_shape = [_shape(x, 0), (_shape(x, 1) - 1) * factor + convH, (_shape(x, 2) - 1) * factor + convW, outC]
        num_groups = _shape(x, 3) // inC

というわけで この strides attributes の値も [0, 2, 3, 1] の順に並び換えたものを用意する。

ここまで準備できたらいよいよ operation の置き換え。

    # New operations
    new_op = graph.create_op(op_type=target_op.type, inputs=inputs, name=f'{target_op.name}_nhwc', attrs=attrs)
    output = tf.transpose(new_op.outputs[0], [0, 3, 1, 2], name=f'{new_op.name}_output')

    # Update connections
    ops = [op for op in graph.get_operations() if target_op.outputs[0] in op.inputs]
    for op in ops:
        for i, input_tensor in enumerate(op.inputs):
            if input_tensor.name == target_op.outputs[0].name:
                op._update_input(i, output)

元の operation と同じ type, attrs を持つ新しい operation を作成する。入力は NHWC に transpose したもの。 ということは この operation の outputs も NHWC になっているので、その後の計算に支障が出ないよう この出力は NCHW に戻しておいてやる必要がある。ので outputs を今度は [0, 3, 1, 2] でtransposeする。

そして、元々の outputs を受け取っていた 次の operation たちの入力を この新しい operation の outputs に置き換えてやる。

PrevOp ---> NCHW ---------> Op(NCHW) -----------NCHW-----> NextOp
        |                                              |
         -> NCHW to NHWC -> Op(NHWC) -> NHWC to NCHW --

元々上段の流れだけだったものに対して、下段のルートを付け足した形になる。

最後に、元々あった NCHW の operation を graph から消しておく。

# Delete old nodes
graph_def = graph.as_graph_def()
for target_op in target_ops:
    graph_def.node.remove(target_op.node_def)

これで、一度 SavedModel に graph を保存してみよう。

inputs = graph.get_tensor_by_name(f'Gs/{Gs.input_names[0]}:0')
outputs = graph.get_tensor_by_name(f'Gs/{Gs.output_names[0]}:0')

tf.compat.v1.enable_resource_variables()
tf.compat.v1.saved_model.simple_save(
    tf.get_default_session(),
    './savedmodel',
    {'inputs': inputs},
    {'outputs': outputs},
)

これを load して実行してみると…

import tensorflow as tf

model = tf.compat.v2.saved_model.load('./savedmodel')
generate = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

z = tf.random.normal([1, 512])
outputs = generate(inputs=z)['outputs']
with tf.compat.v1.Session() as sess:
    sess.run(tf.compat.v1.initializers.tables_initializer())
    print(sess.run(outputs))
2020-02-04 23:24:38.004217: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2020-02-04 23:24:38.014420: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x7fbdbabd65a0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2020-02-04 23:24:38.014438: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
[[[[0.65079993 0.7005807  0.71635807 ... 0.70904994 0.672469
    0.65049875]
   [0.7084702  0.73016286 0.7354313  ... 0.74900943 0.7347464
    0.7381906 ]
   [0.71451813 0.7372785  0.73747885 ... 0.75236607 0.7619566
    0.7417464 ]
   ...

何らかの数値が出力された!

これはやはり NCHW(1, 3, 256, 256) のような形で来ているので、RGBの画像として Pillow などで扱うにはやはり NHWC に transpose したりといった処理は必要になる。

import tensorflow as tf
from PIL import Image

model = tf.compat.v2.saved_model.load('./savedmodel')
generate = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

z = tf.random.normal([1, 512])
outputs = generate(inputs=z)['outputs']
outputs = tf.transpose(outputs, [0, 2, 3, 1])
outputs = tf.saturate_cast((outputs + 1.0) * 127.5, tf.uint8)
with tf.compat.v1.Session() as sess:
    sess.run(tf.compat.v1.initializers.tables_initializer())
    img = Image.fromarray(sess.run(outputs)[0])
img.save('output.png')

f:id:sugyan:20200204233514p:plain

CPU環境でも生成が出来た! やったー!!

2. Modelを構築し直す

標準的な畳み込みを使ったネットワークであれば ここまでの変換で問題なく動くようになるかもしれない。が、StyleGAN2の場合はまだちょっとだけ問題が残っている。

単一の画像生成なら上述のように出来るけど、 batch_size1 より大きくすると、またエラーになってしまう。

z = tf.random.normal([2, 512])
...

tensorflow.python.framework.errors_impl.UnimplementedError: [_Derived_]{{function_node __inference_pruned_14657}} {{function_node __inference_pruned_14657}} The Conv2D op currently does not support grouped convolutions on the CPU. A grouped convolution was attempted to be run because the input depth of 1024 does not match the filter input depth of 512
         [[{{node Gs/G_synthesis/4x4/Conv/Conv2D_nhwc}}]]
         [[StatefulPartitionedCall_1]]

grouped convolution というものが使われていて、これがまた CPU環境ではまだサポートされていないものだった。

このへんかな?

TensorFlow 1.14 あたりから入ったもののようだ。 普通は filter のshapeは [filter_height, filter_width, in_channels, out_channels] と なっていて、inputchannelfilter2 番目の次元と等しくなければならないが、cuDNNによって filter.shape[2] の倍数であればまとめて計算できるようになる、という機能… なのかな。(よく分かっていない)

これがどうやら networks_stylegan2.modulated_conv2d_layer() の中で fused_modconv=True のときにそういった処理をするようになっているらしい。

networks_stylegan2.py 抜粋:

    # Reshape/scale input.
    if fused_modconv:
        x = tf.reshape(x, [1, -1, x.shape[2], x.shape[3]]) # Fused => reshape minibatch to convolution groups.
        w = tf.reshape(tf.transpose(ww, [1, 2, 3, 0, 4]), [ww.shape[1], ww.shape[2], ww.shape[3], -1])
    else:
        x *= tf.cast(s[:, :, np.newaxis, np.newaxis], x.dtype) # [BIhw] Not fused => scale input activations.

    # Convolution with optional up/downsampling.
    if up:
        x = upsample_conv_2d(x, tf.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
    elif down:
        x = conv_downsample_2d(x, tf.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
    else:
        x = tf.nn.conv2d(x, tf.cast(w, x.dtype), data_format='NCHW', strides=[1,1,1,1], padding='SAME')

    # Reshape/scale output.
    if fused_modconv:
        x = tf.reshape(x, [-1, fmaps, x.shape[2], x.shape[3]]) # Fused => reshape convolution groups back to minibatch.

up/downsampling の処理をかける前に xw をreshapeして、その結果を後でまたreshapeして元に戻す、といった形になっているようだ。

これは単一の operation の書き換えでは対応できない…。


ということで、結局 graph の書き換えだけでは無理そうなので 「重みだけ再利用してModelはCPUでも動かせる形に構築し直す」という方針を取ることにした。

まずは .pkl をloadした後、 checkpoint 形式で 変数の値だけを保存する。

import pickle
import tensorflow as tf
from dnnlib import tflib

tflib.init_tf()
with open('network.pkl', 'rb') as fp:
    _, _, Gs = pickle.load(fp, encoding='latin1')

saver = tf.compat.v1.train.Saver(Gs.vars)
saver.save(tf.get_default_session(), './ckpt/network')
$ ls -l ./ckpt
total 240176
-rw-r--r--  1 sugyan  staff         71 Feb  4 23:45 checkpoint
-rw-r--r--  1 sugyan  staff  120838348 Feb  4 23:45 network.data-00000-of-00001
-rw-r--r--  1 sugyan  staff       4898 Feb  4 23:45 network.index
-rw-r--r--  1 sugyan  staff    2114457 Feb  4 23:45 network.meta

で、 fused_modconv=False になるような設定で Generator を作成する。 どうせ graph を構築しなおすことになるのだし、前述したような NCHW -> NHWC の書き換えもコード上でやってしまおう。

Gs/ に関係するところだけなら 以下の3箇所の tf.nn.conv2dtf.nn.conv2d_transpose の入出力まわりを書き換えてやれば大丈夫だ。

diff --git a/dnnlib/tflib/ops/upfirdn_2d.py b/dnnlib/tflib/ops/upfirdn_2d.py
index fd23777..26cc573 100755
--- a/dnnlib/tflib/ops/upfirdn_2d.py
+++ b/dnnlib/tflib/ops/upfirdn_2d.py
@@ -93,7 +93,9 @@ def _upfirdn_2d_ref(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):
     x = tf.transpose(x, [0, 3, 1, 2])
     x = tf.reshape(x, [-1, 1, inH * upy + pady0 + pady1, inW * upx + padx0 + padx1])
     w = tf.constant(k[::-1, ::-1, np.newaxis, np.newaxis], dtype=x.dtype)
-    x = tf.nn.conv2d(x, w, strides=[1,1,1,1], padding='VALID', data_format='NCHW')
+    x = tf.transpose(x, [0, 2, 3, 1])
+    x = tf.nn.conv2d(x, w, strides=[1,1,1,1], padding='VALID', data_format='NHWC')
+    x = tf.transpose(x, [0, 3, 1, 2])
     x = tf.reshape(x, [-1, minorDim, inH * upy + pady0 + pady1 - kernelH + 1, inW * upx + padx0 + padx1 - kernelW + 1])
     x = tf.transpose(x, [0, 2, 3, 1])

@@ -288,7 +290,11 @@ def upsample_conv_2d(x, w, k=None, factor=2, gain=1, data_format='NCHW', impl='c
     w = tf.reshape(w, [convH, convW, -1, num_groups * inC])

     # Execute.
-    x = tf.nn.conv2d_transpose(x, w, output_shape=output_shape, strides=stride, padding='VALID', data_format=data_format)
+    x = tf.transpose(x, [0, 2, 3, 1])
+    stride = [1, factor, factor, 1]
+    output_shape = [_shape(x, 0), (_shape(x, 1) - 1) * factor + convH, (_shape(x, 2) - 1) * factor + convW, outC]
+    x = tf.nn.conv2d_transpose(x, w, output_shape=output_shape, strides=stride, padding='VALID', data_format='NHWC')
+    x = tf.transpose(x, [0, 3, 1, 2])
     return _simple_upfirdn_2d(x, k, pad0=(p+1)//2+factor-1, pad1=p//2+1, data_format=data_format, impl=impl)

 #----------------------------------------------------------------------------
diff --git a/training/networks_stylegan2.py b/training/networks_stylegan2.py
index 6c96fc1..8fe2979 100755
--- a/training/networks_stylegan2.py
+++ b/training/networks_stylegan2.py
@@ -117,7 +117,9 @@ def modulated_conv2d_layer(x, y, fmaps, kernel, up=False, down=False, demodulate
     elif down:
         x = conv_downsample_2d(x, tf.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
     else:
-        x = tf.nn.conv2d(x, tf.cast(w, x.dtype), data_format='NCHW', strides=[1,1,1,1], padding='SAME')
+        x = tf.transpose(x, [0, 2, 3, 1])
+        x = tf.nn.conv2d(x, tf.cast(w, x.dtype), data_format='NHWC', strides=[1,1,1,1], padding='SAME')
+        x = tf.transpose(x, [0, 3, 1, 2])

     # Reshape/scale output.
     if fused_modconv:

で、 Generator を作成し、保存しておいた変数を checkpoint ファイルからloadする。

import tensorflow as tf
from dnnlib import tflib
from dnnlib import EasyDict

tflib.init_tf()
G_args = EasyDict(func_name='training.networks_stylegan2.G_main')
G_args.fused_modconv = False
G = tflib.Network(
    'G',
    num_channels=3,
    resolution=256,
    **G_args)
Gs = G.clone('Gs')

saver = tf.compat.v1.train.Saver(Gs.vars)
saver.restore(tf.get_default_session(), 'ckpt/network')

graph = tf.get_default_graph()
inputs = graph.get_tensor_by_name(f'Gs/{Gs.input_names[0]}:0')
outputs = graph.get_tensor_by_name(f'Gs/{Gs.output_names[0]}:0')
tf.compat.v1.enable_resource_variables()
tf.compat.v1.saved_model.simple_save(
    tf.get_default_session(),
    './savedmodel',
    {'inputs': inputs},
    {'outputs': outputs}
)

これで今度は NCHW の operation も fused_modconv による grouped convolution も含まない SavedModel が出力された、はず。

再度 batch_size > 1 で生成をしてみよう。

model = tf.compat.v2.saved_model.load('./savedmodel')
generate = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

z = tf.random.normal([3, 512])
outputs = generate(inputs=z)['outputs']
outputs = tf.transpose(outputs, [0, 2, 3, 1])
outputs = tf.saturate_cast((outputs + 1.0) * 127.5, tf.uint8)
with tf.compat.v1.Session() as sess:
    sess.run(tf.compat.v1.initializers.tables_initializer())
    img = Image.fromarray(np.concatenate(sess.run(outputs), axis=1))
img.save('output.png')

f:id:sugyan:20200205232314p:plain

ちゃんと3つ画像が生成された! やったー!!

TensorFlow.jsで動かす

さて、ここまで出来ているのならば特殊な operation なども使っていないはずだろうし TensorFlow.js の GraphModel に変換できるだろう、と思うわけです。

$ tensorflowjs_converter \
    --input_format=tf_saved_model \
    --output_format=tfjs_graph_model \
    --signature_name=serving_default \
    --saved_model_tags=serve \
    --skip_op_check \
    ./savedmodel ./tfjs

やはり RandomStandardNormal の operation はサポートされていないので これがある場合は --skip_op_check optionが必要になる。

ちなみに random入力を使っている場所は randomize_noise というパラメータで制御されていて、 Generator の作成時にこれを fused_modconv と同様に G_args.randomize_noise = False と指定しておくとこの operation は使われなくなって、 --skip_op_check でskipする必要が無くなる。

ransomize_noise がどれくらい出力の品質に影響あるかちょっとよく分かっていないけど、問題ないようなら False にしておけば計算の負荷も減るし良いかもしれない。

ともかく、変換した GraphModel を読み込んで実行してみると…

const url = '.../tfjs/model.json'
const randomNormal = (node) => {
    return tf.randomNormal(node.inputs[0].dataSync())
};
tf.registerOp('RandomStandardNormal', randomNormal);
tf.loadGraphModel(url).then((model) => {
    tf.tidy(() => {
        const z = tf.randomNormal([1, 512])
        const results = model.execute(z)
        console.log(results.shape)
    })
}).catch((err) => {
    console.error(err);
})
webgl_util.ts:110 Uncaught Error: Failed to compile fragment shader.

やはり webgl backend ではエラーが出てしまう。 試しに tf.setBackend('cpu') にしてみると、ものすごい時間はかかるが 一応実行は可能なようだ。しかし現実的ではない。

何故 webgl backend では上手くいかないのか… とひたすら地道に探っていたところ、一つ 特殊な箇所を発見した。

upfirdn_2d の reference implementation で、Upsampleするために Rank 6Tensorreshape してから pad を行っている場所がある。

upfirdn_2d.py 抜粋:

def _upfirdn_2d_ref(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):
    """Slow reference implementation of `upfirdn_2d()` using standard TensorFlow ops."""

    ...

    # Upsample (insert zeros).
    x = tf.reshape(x, [-1, inH, 1, inW, 1, minorDim])
    x = tf.pad(x, [[0, 0], [0, 0], [0, upy - 1], [0, 0], [0, upx - 1], [0, 0]])
    x = tf.reshape(x, [-1, inH * upy, inW * upx, minorDim])

    ...

この関数の処理の詳細は理解できていないけど、ともかくこの箇所では 0 で padding することにより Tensor のサイズを大きくしていることだけは分かる。

ところで TensorFlow.js の tf.padAPI reference を見てみると…

Also available are stricter rank-specific methods with the same signature as this method that assert that paddings is of given length.

  • tf.pad1d
  • tf.pad2d
  • tf.pad3d
  • tf.pad4d

…これ、Rank 5 以上のものには対応していないのでは!?

軽く tensorflow/tfjs のコードを見てみたがちょっと分からず… しかしまぁ Rank 4 までしか対応していない、というのは実に有り得る気がする。

ということで 前述の Rank 6Tensor に対する pad の回避するよう処理を書き換えてみることにした。

とはいえ data_format のときのように transpose すれば良いというものでもなさそうだし ちょっとどうすれば良いか分からない…。

が、幸い ここで Upsample するためのパラメータ upy, upx はどうやら 1 もしくは 2 の値でしか渡されてこないらしい、ということが分かった。 1 のときは結局 pad すべき shape は 0 になるので、何もせずに skip してしまえば良い。 2 のときだけ どうにか2箇所だけ shape を増やしてあげる必要がある…。

いや、待てよこれは 1列分だけ 0 padding するだけなのだから、同じ shape の zeros Tensor を後ろから concat してやれば同じ意味になるのでは? と思いついた。

diff --git a/dnnlib/tflib/ops/upfirdn_2d.py b/dnnlib/tflib/ops/upfirdn_2d.py
index fd23777..49d11ed 100755
--- a/dnnlib/tflib/ops/upfirdn_2d.py
+++ b/dnnlib/tflib/ops/upfirdn_2d.py
@@ -82,7 +82,10 @@ def _upfirdn_2d_ref(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):

     # Upsample (insert zeros).
     x = tf.reshape(x, [-1, inH, 1, inW, 1, minorDim])
-    x = tf.pad(x, [[0, 0], [0, 0], [0, upy - 1], [0, 0], [0, upx - 1], [0, 0]])
+    if upy == 2:
+        x = tf.concat([x, tf.zeros(tf.shape(x))], axis=2)
+    if upx == 2:
+        x = tf.concat([x, tf.zeros(tf.shape(x))], axis=4)
     x = tf.reshape(x, [-1, inH * upy, inW * upx, minorDim])

     # Pad (crop if negative).

このように書き換える。

これでまた Generator を作成しなおして SavedModel に保存して GraphModel に変換して…。

f:id:sugyan:20200206005109p:plain

TensorFlow.js でも 動いた!! やったー!!

やはり webgl backend でのエラーは tf.pad が Rank 4 までしか対応していなかったことが原因だったようだ…。いやーまさかこんな方法で解決できるとは。。

実際のところ、 TensorFlow.js での生成を 256x256 サイズで試した感じでは 最初の実行時に 10秒弱かかるが、その後は 数十 ms くらいで計算できてそう。そこから計算結果のデータを取得して、という部分で 3000 ms くらいかかってしまっているが…。

とりあえず 計算はWebWorkerに任せる などして、描画とかUIのところだけ作っていけば、 誰でもブラウザ上で画像生成を試せるようになる、かも…!? という希望は見えた。

ここまでの変更は一応 ここに残しておく。

StyleGAN2による画像生成をJSで動かしたい

memo.sugyan.com

ということで 多少集めることができた画像データを使って、StyleGAN2を使って生成を試してみた。

github.com

こんな感じで学習が進み、ある程度はそれっぽい顔画像が生成できるようになってきた。

f:id:sugyan:20200129231329p:plain

これはまだこれから学習データもさらに増やして改善していくとして、それはそれとして こうした生成をWebブラウザ上でも動かせるようにしたい。 と思った。

とにかくやってみないと分からん、ってことで挑戦。

Pickle to SavedModel

まず、学習の過程で保存されるsnapshot は network-snapshot-000080.pkl のような形で 学習中のネットワーク全体を pickle で固めたものとして保存されている。

これを持ってきて解凍してみるところから。 run_generator.py をみると、以下のように読み込んでいる。

from dnnlib import tflib

tflib.init_tf()
network_pkl = 'network-snapshot-000080.pkl'
with open(network_pkl, 'rb') as fp:
    _G, _D, Gs = pickle.load(fp, encoding='latin1')

しかしこれはGPUが使える環境じゃないと この読み込みさえエラーになってしまう。 dnnlib/tflib/network.py__setstate__ でgraphの構築をしていて、その中で fused_bias_act, upfirdn_2d といったcustom opsを使用している。 これらが、デフォルトでCUDAのAPIを使ったimplementationを使用しているから、のようだ。

もしCPU環境でmodelを読み込もうとしたら、これらをどうにかする必要がある。幸い CUDAを使わない実装も用意されているので

-def fused_bias_act(x, b=None, axis=1, act='linear', alpha=None, gain=None, impl='cuda'):
+def fused_bias_act(x, b=None, axis=1, act='linear', alpha=None, gain=None, impl='ref'):

のように引数のdefault impl'ref' に変えてやることで CUDAが使えない状態でも問題なくloadは出来るようになる。

で、ともかくloadできた Gs から inputs, outputs を取得して SavedModel の形式で保存するようにしてみた。

import pickle
import tensorflow as tf
from dnnlib import tflib


builder = tf.saved_model.Builder('savedmodel')
with tf.Graph().as_default() as graph:
    tflib.init_tf()
    network_pkl = 'network-snapshot-000080.pkl'
    with open(network_pkl, 'rb') as fp:
        _, _, Gs = pickle.load(fp, encoding='latin1')

    inputs = graph.get_tensor_by_name(f'Gs/{Gs.input_names[0]}:0')
    outputs = graph.get_tensor_by_name(f'Gs/{Gs.output_names[0]}:0')
    signature = tf.saved_model.build_signature_def(
        {'inputs': tf.compat.v1.saved_model.build_tensor_info(inputs)},
        {'outputs': tf.compat.v1.saved_model.build_tensor_info(outputs)})
    builder.add_meta_graph_and_variables(
        tf.get_default_session(),
        [tf.saved_model.SERVING],
        signature_def_map={
            tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature
        })
builder.save()

これで、pklのsnapshotからSavedModelの形式に変換することが出来た、ということになる。

$ saved_model_cli show --all --dir savedmodel

MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['inputs'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 512)
        name: Gs/latents_in:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['outputs'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 3, 256, 256)
        name: Gs/images_out:0
  Method name is:

試しにこのSavedModelを読み込んで生成を実行してみる。

import tensorflow as tf
import numpy as np

with tf.Graph().as_default() as graph:
    model = tf.compat.v2.saved_model.load('./savedmodel')

    generate = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
    rnd = np.random.RandomState(0)
    z = tf.constant(rnd.randn(1, 512), tf.float32)

    with tf.compat.v1.Session(graph=graph) as sess:
        sess.run(tf.compat.v1.initializers.global_variables())
        print(sess.run(generate(inputs=z)['outputs']))

すると

2020-01-29 23:33:16.173370: E tensorflow/core/common_runtime/executor.cc:642] Executor failed to create kernel. Invalid argument: Conv2DCustomBackpropInputOp only supports NHWC.
         [[{{node Gs/G_synthesis/8x8/Conv0_up/conv2d_transpose}}]]
Traceback (most recent call last):

...

NHWC でしか動かないものがあって でもmodelは NCHW で作られて固められてしまっているのでダメなようだ…

SavedModel to GraphModel

SavedModelそのままではCPU環境では動かせないのだけど、とりあえず気にせずJSのmodelとしてconvertしてみる。もしかしたらoptimizeの過程でそのへんもどうにかしてくれるかも? と期待してみたり。

pip install tensorflowjs して convertをかけてみる。

$ tensorflowjs_converter \
    --input_format=tf_saved_model \
    --output_format=tfjs_graph_model \
    --signature_name=serving_default \
    --saved_model_tags=serve \
    ./savedmodel \
    ./tfjs

山ほど警告が出た後に、エラーで止まる。

...

ValueError: Unsupported Ops in the model before optimization
RandomStandardNormal

どうやら tensorflowjs_converter では RandomStandardNormal などのopsはサポートされていないらしい。困った。 しかし調べてみると --skip_op_check というoptionが用意されていて、サポートしてないopがあってもとりあえずskipしてくれるようだ。

$ tensorflowjs_converter \
    --input_format=tf_saved_model \
    --output_format=tfjs_graph_model \
    --signature_name=serving_default \
    --saved_model_tags=serve \
    --skip_op_check \
    ./savedmodel \
    ./tfjs

...

2020-01-29 23:43:17.764627: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:786] Optimization results for grappler item: graph_to_optimize
2020-01-29 23:43:17.764712: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   debug_stripper: debug_stripper did nothing. time = 0.969ms.
2020-01-29 23:43:17.764717: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   model_pruner: Graph size after: 2022 nodes (-115), 2099 edges (-115), time = 187.757ms.
2020-01-29 23:43:17.764721: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   constant_folding: Graph size after: 1614 nodes (-408), 1691 edges (-408), time = 1340.328ms.
2020-01-29 23:43:17.764725: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   arithmetic_optimizer: Graph size after: 962 nodes (-652), 1655 edges (-36), time = 209.671ms.
2020-01-29 23:43:17.764728: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   dependency_optimizer: Graph size after: 854 nodes (-108), 1453 edges (-202), time = 108.008ms.
2020-01-29 23:43:17.764805: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   model_pruner: Graph size after: 854 nodes (0), 1453 edges (0), time = 141.52ms.
2020-01-29 23:43:17.764815: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   constant_folding: Graph size after: 854 nodes (0), 1453 edges (0), time = 369.232ms.
2020-01-29 23:43:17.764819: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   arithmetic_optimizer: Graph size after: 854 nodes (0), 1453 edges (0), time = 197.42ms.
2020-01-29 23:43:17.764823: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   dependency_optimizer: Graph size after: 842 nodes (-12), 1441 edges (-12), time = 76.729ms.
2020-01-29 23:43:17.764826: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   debug_stripper: debug_stripper did nothing. time = 10.353ms.
2020-01-29 23:43:17.764829: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   model_pruner: Graph size after: 842 nodes (0), 1441 edges (0), time = 123.947ms.
2020-01-29 23:43:17.764993: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   constant_folding: Graph size after: 842 nodes (0), 1441 edges (0), time = 365.53ms.
2020-01-29 23:43:17.765002: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   arithmetic_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 191.759ms.
2020-01-29 23:43:17.765006: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   dependency_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 77.682ms.
2020-01-29 23:43:17.765010: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   model_pruner: Graph size after: 842 nodes (0), 1441 edges (0), time = 129.848ms.
2020-01-29 23:43:17.765013: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   constant_folding: Graph size after: 842 nodes (0), 1441 edges (0), time = 360.265ms.
2020-01-29 23:43:17.765016: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   arithmetic_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 212.4ms.
2020-01-29 23:43:17.765073: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   dependency_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 79.051ms.
2020-01-29 23:43:22.461213: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:786] Optimization results for grappler item: graph_to_optimize
2020-01-29 23:43:22.461239: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   remapper: Graph size after: 842 nodes (0), 1441 edges (0), time = 149.812ms.
2020-01-29 23:43:22.461244: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   constant_folding: Graph size after: 842 nodes (0), 1441 edges (0), time = 370.785ms.
2020-01-29 23:43:22.461248: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   arithmetic_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 187.354ms.
2020-01-29 23:43:22.461251: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   dependency_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 78.03ms.
2020-01-29 23:43:22.461255: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   remapper: Graph size after: 842 nodes (0), 1441 edges (0), time = 136.545ms.
2020-01-29 23:43:22.461258: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   constant_folding: Graph size after: 842 nodes (0), 1441 edges (0), time = 370.468ms.
2020-01-29 23:43:22.461343: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   arithmetic_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 199.41ms.
2020-01-29 23:43:22.461353: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   dependency_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 79.184ms.
Writing weight file ./tfjs/model.json...

$ ls ./tfjs
group1-shard10of29.bin
group1-shard11of29.bin
group1-shard12of29.bin
group1-shard13of29.bin
group1-shard14of29.bin
group1-shard15of29.bin
group1-shard16of29.bin
group1-shard17of29.bin
group1-shard18of29.bin
group1-shard19of29.bin
group1-shard1of29.bin
group1-shard20of29.bin
group1-shard21of29.bin
group1-shard22of29.bin
group1-shard23of29.bin
group1-shard24of29.bin
group1-shard25of29.bin
group1-shard26of29.bin
group1-shard27of29.bin
group1-shard28of29.bin
group1-shard29of29.bin
group1-shard2of29.bin
group1-shard3of29.bin
group1-shard4of29.bin
group1-shard5of29.bin
group1-shard6of29.bin
group1-shard7of29.bin
group1-shard8of29.bin
group1-shard9of29.bin
model.json

何かが出力された。めでたい。

Load GraphModel and Execute

無事(?)にTensorFlow.jsの GraphModel形式にconvert出来たので、ブラウザ上でこれをloadして実行してみる。

const url = 'http://localhost:8000/tfjs/model.json'
tf.loadGraphModel(url).then((model) => {
    const z = tf.randomUniform([1, 512]);
    model.execute(z).print();
}).catch((err) => {
    console.error(err);
})

するとConsoleで以下のようなエラーが。

 TypeError: Unknown op 'RandomStandardNormal'. File an issue at https://github.com/tensorflow/tfjs/issues so we can add it, or register a custom execution with tf.registerOp()

はいはい。確かにサポートしてないopsだからskipしてもらってたヤツだ。

しかしどうやら tf.registerOp() で自分で定義すれば良いらしい。 tf.randomNormal というAPIはあるので、それがそのまま使えるはず。

const randomNormal = (node) => {
    return tf.randomNormal(node.inputs[0].shape);
};
tf.registerOp('RandomStandardNormal', randomNormal);

というのを足してみる。

そうするとそれっぽく実行できる雰囲気になったのだけど、今度は webgl_util.ts 内で

Error: Failed to compile fragment shader.

という現象が起きて失敗…。えー、、なにそれ。。。

webgl backendでダメなのかな、、と tf.setBackend('cpu') してみると、今度は

Error: Error in conv2dDerInput: depth of input (9) must match input depth for filter 512.

と別のエラー。shapeが合わない…? NCHW のものによる弊害がここで出てきてるのかな…?

現状

ここまでで止まってしまっている状態。JSのmodelをloadするところまでは出来ているから もう一息かな〜と思ったのだけど。。。

modelを一度 NHWC の形式に変換してやるとかの処理が必要になるんだろうか… うーん

顔画像生成のためのデータセットを作る

動機

TensorFlowの登場をきっかけに 機械学習によるアイドル顔識別 という取り組みをしていて、3年以上かけてコツコツとアイドルの自撮りを収集してラベルをつけてデータセットを作ってきたけど、 アイドルヲタクはもう辞めてしまって 現場にも全然行かなくなり、卒業・脱退の情報を追いながらラベルを更新していく作業を続ける情熱はすっかり薄れてしまった。 もうアイドル顔識別プロジェクトは終了にしよう、と思った。

しかし折角今まで集めたデータを捨ててしまうのは勿体無い。せめて最後に何か活用できないものか。 と考えて、「画像生成」に再び取り組んでみることにした。

過去に試したことはあったけど、それほど上手くはいっていない。

この記事を書いたのが2016年。 この後の数年だけでもGANの技術はすさまじく進歩していて、今や 1024x1024 のような高解像度の 写真と見分けがつかないくらいの綺麗な顔画像を生成できるようになったらしい。是非とも試してみたいところだ。

目標

PGGAN (Progressive Growing of GANs) や StyleGAN で使われた CelebA-HQ datasetは、1024x1024 サイズの高解像度の画像を 30,000 枚用意して作られているようだ。

今回はそこまでいかなくとも、せめて 512x512 の画像を 10,000 枚くらいは集めたい。

設計の失敗

しかし自分がアイドル顔識別のために収集してラベル付けしたデータセットは、投稿された自撮り画像から顔領域を検出し 96x96 にリサイズして切り抜いたものだけしか保存していなかった。 あまりストレージに余裕が無くケチった運用をしていたため、元の高解像度の画像をクラウド上に残しておくなどをまったくしていなかった。 つまり 96x96 よりも高解像度の顔画像は手に入らない…。

集め直し

DBから候補となる画像URLを抽出

とはいえ、手元には「元画像のURL」「元画像にひもづいた、抽出された顔画像」「顔画像に対するラベル」のデータは残っている。

  • 各アイドルのTwitterから取得した画像情報 1,654,503
  • 自作の検出器で検出して抽出した顔画像 2,158,681
    • そのうち、人力の手作業でラベル付けしたもの 204,791

などが、自分が3年以上かけて続けたアノテーション作業の成果だ。

高解像度のアイドル顔画像データセットを構築するためには、resize & crop する前の元画像を取得しなおして、今度は解像度を保ったままで顔領域を抽出しなおせば良い。

目当ての「アイドルの自撮り顔画像」だけを選別するには、

  • 写真の中に1枚だけ顔が検出されている
    • → 集合写真などではない単独の自撮りで 高解像度で写っている可能性が高い
  • その顔画像が正しくアイドルとしてラベル付けされている
    • → 顔検出されていても誤検出が一定割合で起きているし、認識対象外のラベル付けをしていたりするので、それらを除外する

という条件のものを抽出すればできるはず。

SELECT
    faces.id,
    photos.id, photos.source_url, photos.photo_url, photos.posted_at,
    labels.id, labels.name
FROM faces
    INNER JOIN photos ON photos.id = faces.photo_id
    INNER JOIN labels ON labels.id = faces.label_id
WHERE
    photos.id in (
        SELECT
            photos.id
        FROM faces
            INNER JOIN photos ON photos.id = faces.photo_id
        WHERE faces.label_id IS NOT NULL
        GROUP BY photos.id
        HAVING COUNT(faces.id) = 1
    )
ORDER BY faces.updated_at DESC

こうして、「おそらくアイドルが単独で写っていたであろう元画像」196,455 枚のURLを取得できた。

しかし 画像URLが取得できていても、それを投稿したアイドルさんが卒業・解散などの後にTwitterアカウントが削除されたり非公開になっていたりすると、もうその画像は参照できなくなってしまう。

実際に取得を試みてダウンロードできたのは このうち 132,513 件だった。

ちょうど休眠アカウント削除というのが最近ニュースになった。卒業後に残っているアイドルのアカウントたちはどうなってしまうのだろうか…。今のうちに画像だけでも取得しておくことが出来て良かったのかもしれない。

Dlibによる単一顔検出

さて、高解像度(といっても 900x1200 程度だけど)の アイドルさんたちの画像を入手することが出来た。

以前はここから OpenCVHaar Feature-based Cascade Classifiers を使って顔検出し、その領域を resize & crop してデータとして使っていた。 また、アイドルの自撮りの特徴として「斜めに傾いて写っているもの」が多く検出しづらい問題があり、それを考慮して回転補正をかけて検出するという仕組みを自作していた。

今回も同様の検出をすることになるが、より高精度に また目・口の位置も検出したいというのもあり、ここでは dlib を使ってみることにした。 dlib は OpenCV同様に顔領域を検出できるほか、その顔領域内のlandmarkとして顔の輪郭や目・鼻・口などの位置まで簡単に検出することができる。

やはり斜めに傾いた顔などにはあまり強くないようなので、以前のものと同様に回転補正をかけて検出を試みるといったことは必要そうだった。 ただ今回はそもそも「対象の画像には顔が一つだけ写っている」という仮定で その単一の顔の部分だけ検出できれば良いので 少し処理は簡単になる。

例えば、宇宙一輝くぴょんぴょこアイドル 宇佐美幸乃ちゃん の場合。

まずは画像を回転することによってはみ出して消えてしまう部分がないように 元画像対角線の長さを持つ正方形領域を作って、その中央に元画像を配置する。

def detect(self, img):
    # Create a large image that does not protrude by rotation
    h, w, c = img.shape
    hypot = math.ceil(math.hypot(h, w))
    hoffset = round((hypot-h)/2)
    woffset = round((hypot-w)/2)
    padded = np.zeros((hypot, hypot, c), np.uint8)
    padded[hoffset:hoffset+h, woffset:woffset+w, :] = img

この画像をそれぞれ少しずつ回転させたものを生成し、それぞれに対して顔検出を試みる。 このとき、 fhog_object_detector.run(image, upsample_num_times, adjust_threshold)APIで検出をかけることで、その検出結果の confidence score も取得できるので それらを含めて全パターンの結果を集める。

手元で試した限りでは -48° 〜 +48° で 12°ずつの回転幅で試すのが、多くの回転角を少ない検出試行で網羅できて良さそうだった。

    self.detector = dlib.get_frontal_face_detector()
    self.predictor = dlib.shape_predictor(datafile)

    ...

    # Attempt detection by rotating at multiple angles
    results = []
    for angle in [-48, -36, -24, -12, 0, 12, 24, 36, 48]:
        rotated = self._rotate(padded, angle)
        dets, scores, indices = self.detector.run(rotated, 0, 0.0)
        if len(dets) == 1:
            results.append([dets[0], scores[0], angle, rotated])
    if len(results) == 0:
        self.logger.info('there are no detected faces')
        return


def _rotate(self, img, angle):
    h, w, _ = img.shape
    mat = cv2.getRotationMatrix2D((w/2, h/2), angle, 1.0)
    return cv2.warpAffine(img, mat, (w, h), cv2.INTER_LANCZOS4)

9つのパターンの中でもっとも高いscoreで顔が検出されたものが、おそらく最も正解に近い傾き角度である、として それを採用する。

f:id:sugyan:20200118212320p:plain

この場合はまったく回転しない 0° でも顔は検出されている(score: 0.3265)が、少し傾けた -12°のものの方が 0.5834 と高いscoreになっているので、そちらを仮の回転角として採用する。

その回転後の画像に対して landmark を検出し、左右の目の中央位置を算出する。 正しく回転して真っ直ぐになっていたら目の高さは同じになるはず、それがズレているのなら そのぶんだけまだ少し傾きがある、という考えで、その左右の目の位置座標から atan2 を使ってその微妙な角度を計算する。

    ...

    # Choose the best angle by scores, and then adjust the angle using the eyes coordinates
    results.sort(key=lambda x: x[1], reverse=True)
    det, _, angle, rotated = results[0]
    shape = self.predictor(rotated, det)
    eyel, eyer = self._eye_center(shape)
    d = eyer - eyel
    angle += math.degrees(math.atan2(d[1], d[0]))
    self.logger.info(f'angle: {angle:.5f}')


    def _eye_center(self, shape):
        eyel, eyer = np.array([0, 0]), np.array([0, 0])
        for i in range(36, 42):
            eyel[0] += shape.part(i).x
            eyel[1] += shape.part(i).y
        for i in range(42, 48):
            eyer[0] += shape.part(i).x
            eyer[1] += shape.part(i).y
        return eyel / 6, eyer / 6

f:id:sugyan:20200118212350p:plain

元々の回転角度と 計算した角度を足して、最終的な回転角とする。 この画像の場合は -12 + 0.156403 = -11.843597° の回転でほぼ真っ直ぐの状態になる、と計算された。

その回転角での画像をもう一度生成し、正しく顔とlandmarkが検出されることを確認する。

    ...

    # Detect face and shapes from adjusted angle
    adjusted = self._rotate(padded, angle)
    dets = self.detector(adjusted)
    if len(dets) != 1:
        self.logger.info('faces are not detected in the rotated image')
        return
    shape = self.predictor(adjusted, dets[0])

次に、見切れている部分の補完を行う。 データセットには顔の周辺部分まで含めて切り取って使うことになるので、その周辺部分で画像が切れていたりすると非常に不自然な領域が存在してしまうことになる。

PGGANの手法では 元の画像から鏡面反射した画像を繋げて広げて(mirror padding)、そこから切り取ることで不自然さを和らげているようだ。 同様にやってみる。

    ...

    # Create a large mirrored image to rotate and crop
    margin = math.ceil(hypot * (math.sqrt(2) - 1.0) / 2)
    mirrored = np.pad(
        img,
        ((hoffset + margin, hypot - h - hoffset + margin),
         (woffset + margin, hypot - w - woffset + margin),
         (0, 0)), mode='symmetric')
    rotated = self._rotate(mirrored, angle)[margin:margin+hypot, margin:margin+hypot, :]

f:id:sugyan:20200118212412p:plain

たしかに背景の壁などはそのまま続いているかのように見えて不自然な領域は減りそうだ。

ここから、両目の位置と口の端の位置・その各点間の距離を使って 切り取るべき顔領域の中心座標と大きさを算出している。 論文内の手法では

  • x: 両目の幅 = e1 - e0
  • y: 両目の中心 から 口の中心 の距離 = (e0 + e1) / 2 - (m0 + m1) / 2
  • c: 切り取る中心座標 = (e0 + e1) / 2 - 0.1 * y
  • s: 切り取るサイズ = max(4.0 * x, 3.6 * y)

といった計算でやっているようだ。そのまま使って適用してみる。

    # Calculate the center position and cropping size
    # https://arxiv.org/pdf/1710.10196v3.pdf
    e0, e1 = self._eye_center(shape)
    m0 = np.array([shape.part(48).x, shape.part(48).y])
    m1 = np.array([shape.part(54).x, shape.part(54).y])
    x = e1 - e0
    y = (e0 + e1) / 2 - (m0 + m1) / 2
    c = (e0 + e1) / 2 + y * 0.1
    s = max(np.linalg.norm(x) * 4.0, np.linalg.norm(y) * 3.6)
    xoffset = int(np.rint(c[0] - s/2))
    yoffset = int(np.rint(c[1] - s/2))
    if xoffset < 0 or yoffset < 0 or xoffset + s >= hypot or yoffset + s >= hypot:
        self.logger.info('cropping area has exceeded the image area')
        return
    size = int(np.rint(s))
    cropped = rotated[yoffset:yoffset+size, xoffset:xoffset+size, :]

f:id:sugyan:20200118212441p:plain

いい感じにそれっぽく、正規化された顔画像として切り抜くことが出来そうだ。

こうして 検出器が出来たので、132,513 件のURLから実際にこの方法による検出を試みた。 そこそこ重い処理ではあるものの、手元のMacBookでも数日かけてゆっくり実行し続けた結果 72,334 件ほどの顔画像を収集することができた。

f:id:sugyan:20200119232129g:plain

見切れ領域の多い画像に起こる問題点

こうして見ると良い画像データが揃っているように見えるが、実際には全然そんなに上手くはいかない。

多くの自撮り画像は かなり寄り気味に撮られていて、顔や頭の輪郭まで全部は写っていない場合が多い。 そうするとどうなるか。鏡面反射で補完しても見切れた顔や頭が反射されて映るだけで 結局不自然な画像になってしまう。

例えば前述の例でも、もしもっと寄り気味に撮られていて頭などが見切れていたら…

f:id:sugyan:20200118214632p:plain

という感じになって、顔やlandmarkは確かに検出されるかもしれないけど、頭や他の部分が変な形に繋がってしまっておかしなものになってしまう。

f:id:sugyan:20200119232246g:plain

ちょっとくらいなら問題ないかもしれないけど、流石に目が複数見えてたりするのはヤバそう…

抽出した顔画像を使った生成テスト

とりあえずは変な形になってしまったデータが存在してしまっていても仕方ない、と割り切って、検出して得ることが出来た顔画像を 10,000 件ほど使って生成モデルでの学習を試みてみた。

512 x 512 にリサイズしたものをデータセットとして使い、 StyleGAN を使って何epochか学習してみた。

f:id:sugyan:20200118215620g:plain

確かにアイドルの顔っぽい画像が生成されるが、やはり右上や左上などに鏡面反射した顔が繋がっているような奇妙な形のものが生成されやすいようだ。。

まぁ、そういう画像を含んだものを学習データとして与えてしまっているのでそうなるのは当然の結果ではある。

画像選別と管理のためのWebアプリケーション

となると今度はデータのクリーニングが必要になってくる。

目視で1枚1枚 画像を確認し、「学習データに使える、顔全体がきれいに入っている画像」と「学習データに使いたくない、不自然な画像」を選別することにした。

ローカル環境でデータを管理したくない、自分好みのUIで作業・確認したい、などの理由もあり、例によって管理用のWebアプリケーションを自作した。

Google App Engine 上で動作するよう、画像を Cloud Storage にアップロード、それにひもづく情報を Cloud Firestore に保存 (以前は Cloud Datastore だったけど 次の時代はFirestoreらしい、ということで今回初めて触ってみた)。Frontendを Create React App で作って、SPAから App Engine Go Standard Environment で作った API を叩く形のアプリケーション。自分しか閲覧・操作できないよう Firebase Authentication で認証するようにしている。

各画像に対して Status というフィールドを用意しておき、 Ready, NG, Pending, OK の4つの値をセットできるようにした。初期値はすべて ReadyReady のものをひたすら見ていって、きれいで使えそうなものを OK、ダメそうなものを NG に変更する。判断に迷ったものは Pending に。1, 2, 3 のボタン操作1つで次々スピーディーに更新していけるようにUIを工夫した。

f:id:sugyan:20200119203714g:plain

誤操作も有り得るので NG だからといって削除したりはせず、 NG として残しておく。これが後で役に立った。

こうして選別作業していって、OK になったものだけを抽出して学習データに使えば、きっときれいな画像が生成できるようになるはず…

選別作業効率化へ

作業は1枚1秒程度でサクサク進むが、実際にやってみると NG の画像が非常に多いことが分かった。

やはり多くのアイドルさんは顔までは写していても頭全体まで写るような自撮りをしていることは少なく、それによってmirror paddingされたものはだいたい頭の形がおかしい画像になってしまう。

8,500 枚ほど選別作業してみてようやく OK のものが 1,250 枚ほど。 約7枚に1枚しか現れず、8割以上は NG もしくは Pending にする感じになった。思った以上に NG の山の中から少数の OK を探すのはストレスフルだし効率が悪い。

NG の 頭の形がおかしくなってるような画像なんて誰でも区別できるし機械にでもやらせればええやん…

と思ったので、一次選別するための機械学習モデルを自作することにした。

幸い、 NG にした画像も削除せずに明確に「NG である」とラベル付けした状態で残している。 画像を入力して、「OKNG か」だけを予測する分類モデルを用意した。 画像に写っている人物が誰か、は関係なく、生成用のデータとして OK なものか NG なものか、だけを判別させるモデルとして学習させることになる。

そこまで厳密に精度を求めるものでもないし、適当に TensorFlow Hub からImageNetで学習済みの InceptionV3 を利用して 2 classes の classification のためのmodelとした。

import tensorflow as tf
import tensorflow_hub as hub

IMAGE_SIZE = (299, 299)


def cnn()
    return hub.KerasLayer("https://tfhub.dev/google/imagenet/inception_v3/feature_vector/4",
                          trainable=trainable, arguments=dict(batch_norm_momentum=0.997))


def train():
    labels = ['ok', 'ng']

    model = tf.keras.Sequential([
        cnn(),
        tf.keras.layers.Dropout(rate=0.1),
        tf.keras.layers.Dense(
            len(labels),
            activation='softmax',
            kernel_regularizer=tf.keras.regularizers.l2(1e-4)),
    ])
    model.build([None, *IMAGE_SIZE, 3])
    model.summary()
    model.compile(
        optimizer=tf.keras.optimizers.RMSprop(),
        loss=tf.keras.losses.CategoricalCrossentropy(),
        metrics=[tf.keras.metrics.CategoricalAccuracy()])

データセットとしては、既にラベル付け済みの NG のものを 4,800 件、 OK のものを 1,200 件 抽出して使用した。 それぞれを 4000:800, 1000:200trainvalidation & test に分割。

最初は全結合層を学習させるだけの転移学習だけでどうにかなるかな、と思ってちょっと試してみたけどダメそうだったので、結局ネットワーク全体を学習させるfine-tuningで。 Google ColaboratoryGPU Runtimeで数十分ほど学習。

NG データの方が多いので学習初期は NG に全振りして accuracy 0.8 とかになるけど、だんだん改善していって 40 epochほど進めると 0.947 まで上がった。

最終的な結果としては、学習に使わなかった validation & test セットに対する推論で 以下のようなConfusion Matrixになった。

f:id:sugyan:20200118235824p:plain

OK label に対しては Precision 0.9016, Recall 0.8250 といったところで、まぁそれなりに学習してくれている、という感覚ではある。

実際に、未使用の Ready の画像 1,000 枚に対してこの分類モデルに判別させてみたところ、197 枚が OK として分類された。 それらを目視で確認してみたところ、 そのうち 128 枚が OK になった。 期待したよりは低かったが、これまでは 数枚に1枚しか現れない OK を探し続ける作業だったのが 今後はこの 一次選別されたものから作業すれば 半数以上は OK を選べるので、作業の心理的ストレスは格段に軽減されて効率的になる。

また、今後さらにデータが増えたら この分類モデルも再度学習させることで、さらに高精度に一次選別を進めることが出来るようになることが期待できる。

現状と今後

こうして、現時点で 1,900 枚くらいまでは OK な画像を集めることができた。 もう少し増えたらそれらを使って生成を再度試してみたいところ。

が、全体の枚数と割合で概算すると 今あるすべての収集画像に対して選別してもまだ OK10,000 枚に届かないかもしれない…。 自撮りのキレイなオススメアイドルさんをご存知の方がいらっしゃったら是非とも教えていただきたいところ。。

Repository

Advent of Code 2019 に挑戦している

f:id:sugyan:20191223224220p:plain

Advent of Code というのがある。

https://adventofcode.com/

日本ではまだあまり 知っている人/やっている人 は多くないかもしれない。検索してみても、日本語の紹介記事はこれくらいしか見つからなかった。

Advent of Code の紹介 - Qiita

僕も、去年 元同僚の @ExAdamu に教えてもらうまでは存在すら知らなかった。

どういうものか、っていうのは上に貼った記事でも書かれている通りで、12/1 〜 12/25 まで 毎日1つずつ、プログラミングを使うパズル問題が出題される、というもの。 puzzle input の入力値が与えられ、それに対する回答を自分の書いたコードで計算し、出力値を submitして正解すれば星が貰える。

入力値とそれに対する正解はどうやらユーザごとに異なるものになっているようで、誰かに正解を訊く みたいなものは出来ないようになっている。 重要なのは正解に辿りつくためのコード、ということになる。

問題は毎日 part1 / part2 と分かれていて、part1はだいたい問題文に書いてある通りに正しく実装すれば答えが出せるような感じになっている。 part2 は、使う入力値はpart1と一緒なのだけど 求められるものがより複雑になり、ちょっと難しくなる。それなりに正しくアルゴリズムとデータ構造を駆使しないと解けなかったり、ある程度は数学的な知識が必要になってきたりするようだ。

すべての日程で part1 / part2 すべて正解すれば星が50個集まる、ということになる。

去年は存在を知っただけで全然挑戦していなかったのだけど、今年はちょっと腕試しと練習を兼ねて、ということで Rustで挑戦してみることにした。 現時点で 23日まで出題されていて、46個中41個まで星を集めた状態。

GitHub - sugyan/adventofcode

別にこれはコンテストとか競技のものではないので 他の人の回答方法を見てはいけないわけではない(早解き上位を目指す人とかは別だろうけど)。

出来る限り自分で考えてみて、分からなかったら reddit で他の人のコードや考え方を見ても良い。 僕も幾つかは詰まってredditの他の人のアイデアを参考にさせてもらったりもした。

とにかく、やってみると、これはとても面白い。

問題がとてもよく出来ていて、スッキリ解けたときの爽快感がすごい。

あと特徴として、毎回恒例なのか今年のが特別なのかは知らないけど シリーズものになっている問題もある。 奇数日は IntCode という 整数値列を使ったレジスタマシンのようなものを使用することになり、このインタプリタを実装する問題が day5, day7, day9 あたりで出されている。このへんは順番にやっていないと解けない。

でも ちゃんと動かせると迷路を出現させたりブロック崩しのゲームが動いたりして、これは感動した

偶数日は逆に他の日とは全然関係ないので まったく予備知識なくても挑戦して解くことが出来るはず。


プログラミングが好きな人や数学が好きな人、是非とも挑戦してみると良いんじゃないかな、と思います。


僕も今年どこまで出来るか分からないけどもうちょっと頑張ってみるし、来年もあったら絶対また挑戦してみたいと思っている。

ISUCON9 本選12位だった

ISUCON9 予選敗退した って書いたんだけど、アレがアレして色々あって、やっぱり本選に出場できることになったので参加してきた。

「失敗から学ぶISUCONの正しい歩き方」チーム、最終結果は12位でした。

言語は予選のときと同じRubyで挑戦。

10:09 初期状態 → 0

初期実装の状態では遅すぎてスコアが出ない、というところからスタート。当然ながらRuby実装に切り替えてもスコアは出ない。

まずはせめてスコアが出るくらいにはしてスタート地点に立たないとね、ってことで諸々サーバ側の準備を進めたりマニュアル読み込んだりしながら 最低限あるべきindex貼って対応していこうね、と作業開始

11:34 ようやくスコア出る → 1,429

さらにindexを貼っていくと、今度はFailするようになる。 実際にブラウザから確認してみると、どうもindexを貼ることによって検索結果の順番が変わってしまうらしい。 primary keyもORDER BY指定もなくて順番が保証されてないやん… 有り得ない、なんだこのクソアプリは!と憤りつつも ORDER BY でそれっぽい順番で返すようにしようとして train_name を指定し、今度はこれが数値に見せかけた文字列だったために 1 の次に 10, 100 が来るというバグを踏み抜き タイムロス。

13:46 ようやくバグが直った → 4,478

一通りindex貼り終わったものが動いて どうにか安心。 そういえば AVAILABLE_DAYS ってのがあるんだよね、ってことで ためしに 1030 に上げてみたところスコア変わらず。 うーん、じゃあ 90 だとどう? → 5,900 まで上がるが500errorなどが出てキツそう。とりあえず 90 で進めることにする。

14:47 複数台構成 → 6,517

id:Soudai さんにお任せしてMySQLをdockerから剥がしてもらって appサーバとdbサーバで分ける構成が動いた。負荷が分散されてスコアが上がる。たまたまこの瞬間で暫定1位の記録になったので慌ててスクショを撮った。

f:id:sugyan:20191007234341p:plain

15:15 get_available_seats 軽減 → 6,672

唯一と言える、僕がまともにコード書いた部分。 /api/train/search 内のN+1はとても複雑で、完全に排除するのは難しそうだったが 最も処理が重くて負担になっているのは get_available_seats を 「premium か否か」「is_smokingか否か」で計4回呼んでいる部分。この中で多数のJOINなどもあるし せめて1回の呼び出しで結果を取得して コード内でカウントすれば多少は早くなるだろう、と id:kamipo さんの提案を受けて実装。

https://github.com/soudai/isucon9-final/pull/5/files

特にバグも無く動いたは動いたが、残念ながらスコアにはあまり寄与しなかった。

そことは別に、なんか benchmarkerの /api/train/search リクエストで 500 errorを多く返していて、ログを見ると departurearrival が検索できず nil になるのにその値を参照しようとしてぬるぽで落ちているものが多数あって、かと言ってそれを応答から消したりエラー握り潰して誤魔化したりしてもどうやってもbenchmarkerに怒られて、どうしようもなくてハマり続けた。これは出題側のバグなんじゃないかなぁとボヤきながら色々試したが結局最後まで解決できず。悔いが残る。

15:39 deadlockの発見

このへんでbenchmarkerがFailしまくるようになり、どうも POST /api/user/reservations/*/cancel で500 errorを頻発しているようだが これも複雑な処理なので 500を返している箇所も複数あり 詳しい原因がよく分からない。

予選のときに得た教訓で、「泥臭いprintデバッグでも活用する」という精神で 画期的なcommitを入れた。

Add 500 error · soudai/isucon9-final@a8afb24 · GitHub

これにより journalctlcancel error!!!!!!!!! 6 を発見し deadlockによるものだというのがすぐに判明し 解決をkamipoさんに丸投げすることが出来た。 これが今回の最大の貢献だったかもしれない……

16:01 deadlock解決(?) → 12,056

これでようやくそこそこのスコアになったが

そこから先 細かいチューニングをするも 結局これより高いスコアが出ることもなく 最後の数十分はひたすら AVAILABLE_DAYS の数値を上げたり下げたりしてFailしないギリギリくらいに…と調整するくらいになってしまった。

しかし 90180くらいで試していて 初期の 10 付近の数値では試そうともしていなかった… もしかしたら敢えてこの数字を下げることでもっと良いスコアが出たかもしれなかったのかなぁ

18:00 競技終了 最終スコア 7,462

10位以内くらいには入りたかったが そこまでもいけなかったなぁ…

反省点

予選のとき以上にコードでの貢献が出来なかった。 1つの処理で何百行もあって複雑なloopのnestがあるような、言ってしまえばクソコード的なものを前に しっかり内容を読み取ってキレイにリファクタリングしていくだけの力が無かった。

どちらかというと方針の相談や 他のメンバーが詰まりそうになったときに横から一緒にみて解決する、みたいな役回りになってしまっていた。それはそれで貢献だったとは思うけど、それだけではどうしても手が足りないし 自分の役目をもっと果たすように動くべきだった、のかもしれない。

POST /api/user/reservations/*/cancel のところは最後まであまり解決できなかったのだけど、外部リクエストの部分を非同期化して まず200を返してしまってから後でゆっくり処理する、などの解法を聞いて、その発想は全然なかったなぁと思った。

感想

同じRubyで挑戦した白金動物園チームが見事に優勝して感動があったし、1人参加の学生さんが2位で驚愕だったし、多くの刺激を受けて とても楽しめたISUCONでした。

今年も本当にありがとうございました!

参照

soudai.hatenablog.com