Tfjs: `StatefulPartitionedCall` is unsupported, preventing OpenAI GPT-2 models from getting converted

[mohataher]

TensorFlow.js version

2.2.0

Browser version

Irrelevant

Describe the problem and code to reproduce

I'm trying to convert a GPT-2 model using tfjs library. Following recommendations on this answer, I exported the GPT-2 model as SavedModel.

!python3 -m pip install -q git+https://github.com/huggingface/transformers.git
!python3 -m pip install tensorflow tensorflowjs

Then ran the following code to export the SavedModel xx.pb file.

from transformers import TFGPT2LMHeadModel, GPT2Tokenizer
import tensorflowjs
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
# add the EOS token as PAD token to avoid warnings
model = TFGPT2LMHeadModel.from_pretrained("gpt2", pad_token_id=tokenizer.eos_token_id)
model.save("./test_gpt2")

Then I ran this command to convert the SavedModel to tfjs compatible file.

!tensorflowjs_converter \
    --input_format=tf_saved_model \
    --output_node_names='gpt2' \
    --saved_model_tags=serve \
    /content/test_gpt2 \
    /content/test_gpt2_web_model

This causes an error

2020-07-08 16:36:11.455383: I tensorflow/core/platform/cpu_feature_guard.cc:143] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2020-07-08 16:36:11.459979: I tensorflow/core/platform/profile_utils/cpu_utils.cc:102] CPU Frequency: 2300000000 Hz
2020-07-08 16:36:11.460216: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x2e5b100 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2020-07-08 16:36:11.460284: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
2020-07-08 16:36:18.337463: I tensorflow/core/grappler/devices.cc:60] Number of eligible GPUs (core count >= 8, compute capability >= 0.0): 0 (Note: TensorFlow was not compiled with CUDA support)
2020-07-08 16:36:18.337631: I tensorflow/core/grappler/clusters/single_machine.cc:356] Starting new session
2020-07-08 16:36:18.536301: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:797] Optimization results for grappler item: graph_to_optimize
2020-07-08 16:36:18.536373: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:799]   function_optimizer: Graph size after: 163 nodes (0), 175 edges (0), time = 43.871ms.
2020-07-08 16:36:18.536384: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:799]   function_optimizer: Graph size after: 163 nodes (0), 175 edges (0), time = 50.779ms.
2020-07-08 16:36:18.536393: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:797] Optimization results for grappler item: __inference__wrapped_model_24863
2020-07-08 16:36:18.536402: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:799]   function_optimizer: function_optimizer did nothing. time = 0.004ms.
2020-07-08 16:36:18.536411: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:799]   function_optimizer: function_optimizer did nothing. time = 0ms.
Traceback (most recent call last):
  File "/usr/local/bin/tensorflowjs_converter", line 8, in <module>
    sys.exit(pip_main())
  File "/usr/local/lib/python3.6/dist-packages/tensorflowjs/converters/converter.py", line 735, in pip_main
    main([' '.join(sys.argv[1:])])
  File "/usr/local/lib/python3.6/dist-packages/tensorflowjs/converters/converter.py", line 739, in main
    convert(argv[0].split(' '))
  File "/usr/local/lib/python3.6/dist-packages/tensorflowjs/converters/converter.py", line 681, in convert
    control_flow_v2=args.control_flow_v2)
  File "/usr/local/lib/python3.6/dist-packages/tensorflowjs/converters/tf_saved_model_conversion_v2.py", line 494, in convert_tf_saved_model
    weight_shard_size_bytes=weight_shard_size_bytes)
  File "/usr/local/lib/python3.6/dist-packages/tensorflowjs/converters/tf_saved_model_conversion_v2.py", line 143, in optimize_graph
    ', '.join(unsupported))
ValueError: Unsupported Ops in the model before optimization
StatefulPartitionedCall

It says StatefulPartitionedCall is unsupported. Could it be supported please?

[rthadur]

@mohataher this has been supported in latest release 2.0.1 , please verify

[mohataher]

Heya, I can confirm the StatefulPartitionedCall is supported. However, I'm still getting errors related to the same class during the conversion process.

Please see this Google Colab to regenerate the error. https://colab.research.google.com/drive/1YNEnnNxJDjnr-rxun0yVSdUmPFdsgbY-?usp=sharing

The error is copied here for reference InvalidArgumentError: Input 1 of node StatefulPartitionedCall was passed float from unknown:0 incompatible with expected resource.

Traceback (most recent call last):
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/importer.py", line 497, in _import_graph_def_internal
    graph._c_graph, serialized, options)  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.InvalidArgumentError: Input 1 of node StatefulPartitionedCall was passed float from unknown:0 incompatible with expected resource.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/lib/python3.6/dist-packages/tensorflowjs/converters/tf_saved_model_conversion_v2.py", line 482, in convert_tf_saved_model
    frozen_graph = _freeze_saved_model_v2(concrete_func, control_flow_v2)
  File "/usr/local/lib/python3.6/dist-packages/tensorflowjs/converters/tf_saved_model_conversion_v2.py", line 352, in _freeze_saved_model_v2
    concrete_func, lower_control_flow=not control_flow_v2).graph
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/convert_to_constants.py", line 1075, in convert_variables_to_constants_v2
    converted_input_indices)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/convert_to_constants.py", line 1001, in _construct_concrete_function
    new_output_names)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/wrap_function.py", line 650, in function_from_graph_def
    wrapped_import = wrap_function(_imports_graph_def, [])
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/wrap_function.py", line 628, in wrap_function
    collections={}),
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/func_graph.py", line 986, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/wrap_function.py", line 87, in __call__
    return self.call_with_variable_creator_scope(self._fn)(*args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/wrap_function.py", line 93, in wrapped
    return fn(*args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/wrap_function.py", line 648, in _imports_graph_def
    importer.import_graph_def(graph_def, name="")
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/importer.py", line 405, in import_graph_def
    producer_op_list=producer_op_list)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/importer.py", line 501, in _import_graph_def_internal
    raise ValueError(str(e))
ValueError: Input 1 of node StatefulPartitionedCall was passed float from unknown:0 incompatible with expected resource.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/bin/tensorflowjs_converter", line 8, in <module>
    sys.exit(pip_main())
  File "/usr/local/lib/python3.6/dist-packages/tensorflowjs/converters/converter.py", line 735, in pip_main
    main([' '.join(sys.argv[1:])])
  File "/usr/local/lib/python3.6/dist-packages/tensorflowjs/converters/converter.py", line 739, in main
    convert(argv[0].split(' '))
  File "/usr/local/lib/python3.6/dist-packages/tensorflowjs/converters/converter.py", line 681, in convert
    control_flow_v2=args.control_flow_v2)
  File "/usr/local/lib/python3.6/dist-packages/tensorflowjs/converters/tf_saved_model_conversion_v2.py", line 485, in convert_tf_saved_model
    output_node_names)
  File "/usr/local/lib/python3.6/dist-packages/tensorflowjs/converters/tf_saved_model_conversion_v2.py", line 342, in _freeze_saved_model_v1
    sess, g.as_graph_def(), output_node_names)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/deprecation.py", line 324, in new_func
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/graph_util_impl.py", line 280, in convert_variables_to_constants
    variable_names_blacklist=variable_names_blacklist)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/convert_to_constants.py", line 1146, in convert_variables_to_constants_from_session_graph
    variable_names_blacklist=variable_names_blacklist))
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/convert_to_constants.py", line 854, in __init__
    graph_def = graph_util.extract_sub_graph(graph_def, output_node_names)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/deprecation.py", line 324, in new_func
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/graph_util_impl.py", line 208, in extract_sub_graph
    _assert_nodes_are_present(name_to_node, dest_nodes)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/graph_util_impl.py", line 163, in _assert_nodes_are_present
    assert d in name_to_node, "%s is not in graph" % d
AssertionError: Identity is not in graph

I'm not sure if this is still related to the current issue or another issue should be opened. It's up to you. But I'm still unable to convert a GPT-2 model to tfjs.

[pyu10055]

@mohataher I am having trouble with the colab, can you try to add following flag to the conversion

--control_flow_v2=true

[mohataher]

I have used the standard tensorflow version 2.2.0 ... Same error above.

When using the tf-nightly, I'm getting an entirely different error but that's probably out of the scope of this issue.

[VogtAI]

Any progress on this? Also trying to convert GPT-2 to tfjs and failing...

[EddieOne]

+1 Anyone try this? https://stackoverflow.com/a/62679227/654499

[VogtAI]

+1 Anyone try this? https://stackoverflow.com/a/62679227/654499

That's the same code mohataher used.

[elbowdonkey]

I see the same StatefulPartitionedCall related error as @mohataher when trying to run tensorflowjs_converter. My arguments are the same, with the exception of what the converted models should be named.

I have some version differences that might be useful to know:

• Python: 3.8.3
• tensorflow-cpu: 2.3.0
• tensorflowjs: 2.0.1.post1

[elbowdonkey]

@EddieOne and yeah, that SO link is the exact code that @mohataher and I have used to produce the error.

[elbowdonkey]

@pyu10055 adding --control_flow_v2=true to the tensorflowjs_converter command doesn't change the StatefulPartitionedCall related error.

[pyu10055]

@elbowdonkey @mohataher The issue is the model contains stateful variables, (i.e. stateful keras LSTM layer), we cannot currently freeze the graph that contains stateful variable. One way to work around is to use stateless layers.

[pyu10055]

@elbowdonkey @mohataher This is likely fixed by this PR https://github.com/tensorflow/tfjs/pull/3685
You can try again one 2.1.0 is released, should released this week or early next week.

[pyu10055]

@elbowdonkey @mohataher I have verified gpt2 model is convertible after PR #3685, please wait for 2.1.0 release. thanks

[VogtAI]

@elbowdonkey @mohataher I have verified gpt2 model is convertible after PR #3685, please wait for 2.1.0 release. thanks

It works now. Thanks!

[tensorflow-butler[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed in 7 dyas if no further activity occurs. Thank you.

[tensorflow-butler[bot]]

Closing as stale. Please @mention us if this needs more attention.

[tensorflow-butler[bot]]

Are you satisfied with the resolution of your issue?
Yes
No