search

tensorflow / tfjs

A WebGL accelerated JavaScript library for training and deploying ML models.
https://js.tensorflow.org
Apache License 2.0
18.51k stars 1.93k forks source link

Trying to port nvidia stylegan model to tensorflowjs, need help with operator Assign, IdentityN, VariableV2 #1296

Closed yjhatfdu closed 5 years ago

yjhatfdu commented 5 years ago

TensorFlow.js version

0.15.3

Browser version

Chrome 72

Describe the problem or feature request

I tried to port Nvidia stylegan pretrained model to tfjs, need help with operator Assign, IdentityN, VariableV2

Code to reproduce the bug / link to feature request

I used the flowing code to dump stylegan pkl model to tensorflow saved model.

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

url = 'https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ'  # karras2019stylegan-ffhq-1024x1024.pkl
     with dnnlib.util.open_url(url, cache_dir=config.cache_dir) as f:
     _G, _D, Gs = pickle.load(f)
Gs.print_layers()
sess = tf.get_default_session()
builder = tf.saved_model.builder.SavedModelBuilder('output_model')
builder.add_meta_graph_and_variables(sess, [tf.saved_model.tag_constants.SERVING], strip_default_attrs=True)
builder.save()

Then

tensorflowjs_converter --skip_op_check=SKIP_OP_CHECK  --input_format=tf_saved_model --saved_model_tags=serve --output_node_names='Gs/images_out' output_model web

Result is

Traceback (most recent call last):
  File "/usr/local/bin/tensorflowjs_converter", line 10, in <module>
    sys.exit(main())
  File "/usr/local/lib/python3.7/site-packages/tensorflowjs/converters/converter.py", line 322, in main
    strip_debug_ops=FLAGS.strip_debug_ops)
  File "/usr/local/lib/python3.7/site-packages/tensorflowjs/converters/tf_saved_model_conversion_pb.py", line 291, in convert_tf_saved_model
    graph = load_graph(output_graph + '.frozen', output_node_names)
  File "/usr/local/lib/python3.7/site-packages/tensorflowjs/converters/tf_saved_model_conversion_pb.py", line 67, in load_graph
    tf.import_graph_def(graph_def, name='')
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/framework/importer.py", line 431, in import_graph_def
    raise ValueError(str(e))
ValueError: Input 0 of node Gs/Assign was passed float from G_synthesis_1/lod:0 incompatible with expected float_ref.

I added 'G_synthesis_1/lod:0' to freeze_graph variable_names_blacklist to bypass this error. At last, load this model on browser.

async load(){
    const MODEL_URL = '/assets/web/tensorflowjs_model.pb';
    const WEIGHTS_URL = '/assets/web/weights_manifest.json';
    const model = await tf.loadFrozenModel(MODEL_URL, WEIGHTS_URL);
  }

We need Assign, IdentityN, VariableV2 now.

yjhatfdu commented 5 years ago

I modified the model and code, there are many more ops mode and arguments not supported

knok commented 5 years ago

I am trying same task with tensorflowjs 1.2.3, but still have same error. The follwing is my dumper program:

import pickle
import numpy as np
import dnnlib.tflib as tflib
import dnnlib
import tensorflow as tf
import config

tflib.init_tf()
url = 'https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ'  # my trained data
with dnnlib.util.open_url(url, cache_dir=config.cache_dir) as f:
    _G, _D, Gs = pickle.load(f)
Gs.print_layers()
layers = Gs.list_layers()
latents_in = layers[0][1]
images_out = layers[-1][1]
sess = tf.get_default_session()
builder = tf.saved_model.builder.SavedModelBuilder('output_model')
builder.add_meta_graph_and_variables(sess, \
    [tf.saved_model.tag_constants.SERVING], \
        signature_def_map = {
            'serving_default': tf.saved_model.signature_def_utils.build_signature_def(
                inputs={'latents_in': tf.saved_model.utils.build_tensor_info(latents_in)},
                outputs={'images_out':  tf.saved_model.utils.build_tensor_info(images_out)},
                method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME,
            ),
        },
    strip_default_attrs=True)
builder.save()

Then run converter:

$ tensorflowjs_converter --input_format tf_saved_model \
--output_format tfjs_graph_model \
output_model web_model --skip_op_check

It returns error with the following trackback:

Traceback (most recent call last):
  File "/opt/p3/lib/python3.7/site-packages/tensorflow/python/framework/importer.py", line 427, in import_graph_def
    graph._c_graph, serialized, options)  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.InvalidArgumentError: Input 0 of node Gs/Assign was passed float from G_synthesis_1/lod:0 incompatible with expected float_ref.

I don't know how to avoid the error with variable_names_blacklist.

knok commented 5 years ago

I'm trying to invoke freeze_graph. Save checkpoint (modify generate_figure.py) and invoke freeze_graph:

$ freeze_graph --input_graph output_model/saved_model.pb \
 --input_checkpoint checkpoint/model.ckpt \
 --input_binary=true --output_graph output_model/fix.pb \
 --output_node_names images_out

Traceback (most recent call last):
  File "/home/ubuntu/anaconda3/envs/tensorflow_p36/bin/freeze_graph", line 11, in <module>
    sys.exit(run_main())
  File "/home/ubuntu/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow/python/tools/freeze_graph.py", line 488, in run_main
    app.run(main=my_main, argv=[sys.argv[0]] + unparsed)
  File "/home/ubuntu/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow/python/platform/app.py", line 125, in run
    _sys.exit(main(argv))
  File "/home/ubuntu/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow/python/tools/freeze_graph.py", line 487, in <lambda>
    my_main = lambda unused_args: main(unused_args, flags)
  File "/home/ubuntu/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow/python/tools/freeze_graph.py", line 381, in main
    flags.saved_model_tags, checkpoint_version)
  File "/home/ubuntu/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow/python/tools/freeze_graph.py", line 340, in freeze_graph
    input_graph_def = _parse_input_graph_proto(input_graph, input_binary)
  File "/home/ubuntu/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow/python/tools/freeze_graph.py", line 251, in _parse_input_graph_proto
    input_graph_def.ParseFromString(f.read())
google.protobuf.message.DecodeError: Error parsing message

It seems another problem (and maybe my fault).

jochemstoel commented 4 years ago

Up. Has anyone succeeded at this?

saadz-khan commented 4 years ago

Hey @yjhatfdu @knok does the code you guys shared above without tensorflow.js converter scripts run works and convert the stylegan model to tensorflow format? What is the import config as it is giving an error for me? Any specific version or where is this config imported from??

import pickle
import numpy as np
import dnnlib.tflib as tflib
import dnnlib
import tensorflow as tf
import config

tflib.init_tf()
url = 'https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ'  # my trained data
with dnnlib.util.open_url(url, cache_dir=config.cache_dir) as f:
    _G, _D, Gs = pickle.load(f)
Gs.print_layers()
layers = Gs.list_layers()
latents_in = layers[0][1]
images_out = layers[-1][1]
sess = tf.get_default_session()
builder = tf.saved_model.builder.SavedModelBuilder('output_model')
builder.add_meta_graph_and_variables(sess, \
    [tf.saved_model.tag_constants.SERVING], \
        signature_def_map = {
            'serving_default': tf.saved_model.signature_def_utils.build_signature_def(
                inputs={'latents_in': tf.saved_model.utils.build_tensor_info(latents_in)},
                outputs={'images_out':  tf.saved_model.utils.build_tensor_info(images_out)},
                method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME,
            ),
        },
    strip_default_attrs=True)
builder.save()