Hub: USE latest version can't be used with MirroredStrategy ("Trying to access a placeholder that is not supposed to be executed.")

[eduardofv]

The latest versions of USE throw this error when used in a MirroredStrategy. (tf: 2.1.0, keras: 2.2.4-tf, hub: 0.7.0)

import tensorflow as tf
import tensorflow.keras as keras
import tensorflow_hub as hub
import tensorflow_text as text

#This USE models fail with " InvalidArgumentError:  assertion failed: [Trying to access a placeholder that is not supposed to be executed"
LM = "https://tfhub.dev/google/universal-sentence-encoder-multilingual-large/3"
#LM = "https://tfhub.dev/google/universal-sentence-encoder-multilingual/3"
#LM = "https://tfhub.dev/google/universal-sentence-encoder-large/5"
#LM = "https://tfhub.dev/google/universal-sentence-encoder/4"
DIM = 512

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    model = keras.models.Sequential()
    model.add(
        hub.KerasLayer(LM,
                       output_shape=DIM,
                       input_shape=[],
                       dtype=tf.string)
    )
    model.add(keras.layers.Dense(1, activation='sigmoid'))
    model.compile(optimizer="adam", loss="binary_crossentropy")

model.summary()

Throws:

INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-2-6b41a257fb9c> in <module>()
     32                        output_shape=DIM,
     33                        input_shape=[],
---> 34                        dtype=tf.string)
     35     )
     36     model.add(keras.layers.Dense(1, activation='sigmoid'))

12 frames
/usr/local/lib/python3.6/dist-packages/six.py in raise_from(value, from_value)

InvalidArgumentError:  assertion failed: [Trying to access a placeholder that is not supposed to be executed. This means you are executing a graph generated from cross-replica context in an in-replica context.]
     [[node Assert/Assert (defined at /usr/local/lib/python3.6/dist-packages/tensorflow_hub/module_v2.py:95) ]] [Op:__inference_restored_function_body_45855]

Function call stack:
restored_function_body

Interestingly, NNLM works fine and also USE-large v3 (albeit some warnings, not sure if they affect it's later performance). Also, if OneDeviceStrategy or no strategy is used everything works.

Check this colab with code and test cases: https://colab.research.google.com/drive/1_YaGYje4tXPyQDx_hYaj9VNLAA5hi3Dg

[mercarikaicheung]

Any updates on this also facing this issue?

[arnoegw]

Confirmed: there is an issue. Thank you, @eduardofv, for the very clear and reproducible report!

@mercarikaicheung, there is no useful update yet, sorry.

[guptapriya]

@arnoegw looking at the stacktrace, it seems that the model.add(hub.KerasLayer(...)) is actually trying to execute some ops. Is that expected? Usually I would expect loading a model to only construct the model but not try to execute it. Also, do you know what the differences between USE version 5 and 3 are (since 3 seems to work but not 5?)

[eduardofv]

From the documentation of USE large: https://tfhub.dev/google/universal-sentence-encoder-large/5

Changelog

Version 1

• Initial release.

Version 2

• Exposed internal variables as Trainable.

Version 3

• Fixed batch invariant bug. This version was retrained and its embedding space differs from previous versions.

Version 4

• Retrained using TF2.

Version 5

• Fixed GitHub issue #409.
• Default inference function now returns the Tensor instead of a dictionary.

[mercarikaicheung]

Hi everyone,
I was wondering if there are any updates on this?

[guptapriya]

Not from the distribution strategy side since it appears that this model might have been saved in a non standard way. The TF Hub team is working on new version of the USE that will be built natively in TF2 and should work correctly.

[hi-im-ryanli]

I am running into the same issue. Any timeline on the fixes, @guptapriya ?

[arnoegw]

@jaxlaw?

[crccw]

Hi all, the issue was because this model was converted from TF1 SavedModel to TF2 SavedModel. There's some tricky issues with loading such models under tf.distribute.Strategy. The team is working on a TF2 native version of the model.

[r-wheeler]

I am able to reproduce this same error using bert on 1 machine with 8 v100 gpus using mirrored strategy. It did not occur when running the same code (and data) on 4 v100 gpus

Traceback (most recent call last):
...
tensorflow.python.framework.errors_impl.InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument:  assertion failed: [Trying to access a placeholder that is not supposed to be executed. This means you are executing a graph generated from the cross-replica context in an in-replica context.]
     [[{{node bert_embeddings/keras_layer/StatefulPartitionedCall/Assert/Assert}}]]
     [[cond/else/_144/Maximum/_388]]
  (1) Invalid argument:  assertion failed: [Trying to access a placeholder that is not supposed to be executed. This means you are executing a graph generated from the cross-replica context in an in-replica context.]
     [[{{node bert_embeddings/keras_layer/StatefulPartitionedCall/Assert/Assert}}]]
0 successful operations.
0 derived errors ignored. [Op:__inference_call_240207]

Function call stack:
call -> call

[RobRomijnders]

Changing from strategy.scope to strategy.run resolved the issue for me. Not sure if this solution is related to the thread though.

[eduardofv]

Changing from strategy.scope to strategy.run resolved the issue for me. Not sure if this solution is related to the thread though.

Great. Currently I'm not working on this but if I return to it I'll check this solution. Thanks

[dkorkinof]

Hi all, any news on this issue?

[RobRomijnders]

@dkorkinof What issue are you facing exactly? Irrespective of the distribution strategy, defining ops within strategy.scope and accessing them via strategy.run resolved the issue for me.

[dkorkinof]

Basically, I am getting the exact same error with @eduardofv.
More specifically with TF 2.3.1 and Hub 0.10.0 the following code:

import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_text

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    model = hub.load("https://tfhub.dev/google/universal-sentence-encoder-multilingual-large/3")

Throws this exception:

Traceback (most recent call last):
  File "test.py", line 10, in <module>
    model = hub.load("https://tfhub.dev/google/universal-sentence-encoder-multilingual-large/3")
  File "python3.7/site-packages/tensorflow_hub/module_v2.py", line 106, in load
    obj = tf.compat.v1.saved_model.load_v2(module_path, tags=tags)
  File "python3.7/site-packages/tensorflow/python/saved_model/load.py", line 603, in load
    return load_internal(export_dir, tags, options)
  File "python3.7/site-packages/tensorflow/python/saved_model/load.py", line 633, in load_internal
    ckpt_options)
  File "python3.7/site-packages/tensorflow/python/saved_model/load.py", line 135, in __init__
    init_op = node._initialize()  # pylint: disable=protected-access
  File "python3.7/site-packages/tensorflow/python/eager/def_function.py", line 780, in __call__
    result = self._call(*args, **kwds)
  File "python3.7/site-packages/tensorflow/python/eager/def_function.py", line 846, in _call
    return self._concrete_stateful_fn._filtered_call(canon_args, canon_kwds)  # pylint: disable=protected-access
  File "python3.7/site-packages/tensorflow/python/eager/function.py", line 1848, in _filtered_call
    cancellation_manager=cancellation_manager)
  File "python3.7/site-packages/tensorflow/python/eager/function.py", line 1924, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager))
  File "python3.7/site-packages/tensorflow/python/eager/function.py", line 550, in call
    ctx=ctx)
  File "python3.7/site-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.InvalidArgumentError:  assertion failed: [Trying to access a placeholder that is not supposed to be executed. This means you are executing a graph generated from the cross-replica context in an in-replica context.]
     [[node Assert/Assert (defined at python3.7/site-packages/tensorflow_hub/module_v2.py:106) ]] [Op:__inference_restored_function_body_49292]

Function call stack:
restored_function_body

It seems what @crccw mentioned checks out, as the model's tensorflow version appears as 1.15.0.
However automatically converting the TF1 model to TF2 is different to what is mentioned in the release notes: "Version 4 Retrained using TF2.".
Has there been any progress since @crccw 's last post in July?