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