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

memo.sugyan.com

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

StyleGAN2 がゼロから学習してアイドルの顔画像っぽいものを習得していく過程、こんな感じ pic.twitter.com/jTofmFBlxF
— すぎゃーん💯 (@sugyan) January 27, 2020

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

f:id:sugyan:20200129231329p:plain

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

StyleGANのモデルをTensorFlow.jsなどにconvertしてWebブラウザ上で動かす、というのは出来ないのかな…？軽く調べた感じではまだ成功例が無さそう
— すぎゃーん💯 (@sugyan) January 27, 2020

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

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 の形式に変換してやるとかの処理が必要になるんだろうか… うーん