ValueError: Unsupported Ops in the model before optimization StatefulPartitionedCall

[msahamed]

TensorFlow.js version: 2.0.1.post1

Browser version: Chrome Version 83.0.4103.106 (Official Build) (64-bit)

Describe the problem or feature request

I am trying to convert a TensorFlow saved_model to TensorFlowJS type with a "quantize" flag. The model uses Huggingface's Distillbert as a layer as follows.

config = DistilBertConfig.from_pretrained( 'distilbert-base-uncased')
config.output_hidden_states = False
distillbert_main = TFDistilBertMainLayer(config = config)

input_word_ids = tf.keras.layers.Input(shape=(8,), dtype = tf.int32, name = "input_word_ids"),
x = distillbert_main(input_word_ids)[0]
x = tf.keras.layers.Lambda(lambda seq: seq[:, 0, :])(x)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Dropout(0.2)(x)
out = tf.keras.layers.Dense(2)(x)

model = tf.keras.Model(inputs=input_word_ids, outputs=out)
for layer in model.layers[:3]:
    layer.trainable = False
model.summary() # Works fine
model.get_config() # Works fine
tf.saved_model.save(model, './models/model') # works fine
tfjs.converters.convert_tf_saved_model('./models/model', './models/tfjs') # does not work

I know tfjs currently supports Image-based models, and Models with control flow ops (e.g., RNNs) are also supported. It would be great to see if tfjs team could add the features that convert the attention-based models like the model above.

Code to reproduce the bug / link to feature request

tensorflowjs_converter --quantize_float16 --input_format=tf_saved_model /Users/sabber/Desktop/node_test/model /Users/sabber/Desktop/node_test/model/quantized

[rthadur]

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

[msahamed]

Thanks for the reply @rthadur. I updated TensorFlow.js from 2.0.1.post1 to tensorflowjs 2.0.1 but still the same error:

ValueError: Unsupported Ops in the model before optimization
StatefulPartitionedCall

[rthadur]

did you try to create a VM to convert ?

[msahamed]

did you try to create a VM to convert ?

Nope, I have not, in a virtual environment? I am confused.

[rthadur]

Here are the steps to convert using a VM.

[msahamed]

Here are the steps to convert using a VM.

The link takes to a different link (steps link ) that requires authentication. Could you please help me without creating VM but using tensorflowjs b like below:

tfjs.converters.convert_tf_saved_model('./models/model', './models/tfjs')

steps link: https://login.corp.google.com/request?s=copybara-safe-internal-review.git.corp.google.com:443/uberproxy/&d=https://copybara-safe-internal-review.git.corp.google.com/q/59600%3Fupxsrf%3DAO9zFu0-dZnj-jl3EX9IjvVz-r8HfxLEOQBV7ulxKGlkefT0SQ:1594866494972&maxAge=1200&authLevel=2000000&keyIds=zRI,s8x&3pd=1

[pyu10055]

@msahamed Thank you for reporting, can you provide the full python code that I can reproduce the problem. Also can you provide the pips that are needed, thanks.

[msahamed]

@msahamed Thank you for reporting, can you provide the full python code that I can reproduce the problem. Also can you provide the pips that are needed, thanks.

Thank you @pyu10055 for your reply. To reproduce you need tensorflow (2.2.0), tensorflowjs (2.0.1) and transformers (3.0.0). Below are the codes and errors:

Code:

import tensorflow as tf # 2.2.0
import tensorflowjs as tfjs # 2.0.1
from transformers import DistilBertConfig, TFDistilBertModel # 3.0.0

config = DistilBertConfig()
config.output_hidden_states = False
distil_bert_model = TFDistilBertModel.from_pretrained('distilbert-base-uncased', config = config)

input_word_ids = tf.keras.layers.Input(shape=(8,), dtype = tf.int32, name = "input_word_ids"),
x = distil_bert_model(input_word_ids)[0]
x = tf.keras.layers.Lambda(lambda seq: seq[:, 0, :])(x)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Dense(256, activation='relu')(x)
x = tf.keras.layers.Dropout(0.2)(x)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Dense(128, activation='relu')(x)
x = tf.keras.layers.Dropout(0.2)(x)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Dense(64, activation='relu')(x)
x = tf.keras.layers.Dropout(0.2)(x)
out = tf.keras.layers.Dense(2)(x)

model = tf.keras.Model(inputs=input_word_ids, outputs=out)

for layer in model.layers[:3]:
    layer.trainable = False

model.summary()

model.compile(
    loss = 'mse', 
    optimizer = tf.keras.optimizers.Adam(0.0005)
)

tf.saved_model.save(model, './models/model')
model.save_weights('./models/model/model_weight.h5') # works everything fine
tfjs.converters.convert_tf_saved_model('./models/model', './models/tfjs') # does not work

Error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-12-354a49fe97b3> in <module>
----> 1 tfjs.converters.convert_tf_saved_model('./models/model', './models/tfjs')

/usr/local/lib/python3.7/site-packages/tensorflowjs/converters/tf_saved_model_conversion_v2.py in convert_tf_saved_model(saved_model_dir, output_dir, signature_def, saved_model_tags, quantization_dtype_map, skip_op_check, strip_debug_ops, weight_shard_size_bytes, control_flow_v2)
    492                  skip_op_check=skip_op_check,
    493                  strip_debug_ops=strip_debug_ops,
--> 494                  weight_shard_size_bytes=weight_shard_size_bytes)
    495 
    496 def load_and_initialize_hub_module(module_path, signature='default'):

/usr/local/lib/python3.7/site-packages/tensorflowjs/converters/tf_saved_model_conversion_v2.py in optimize_graph(graph, signature_def, output_graph, tf_version, quantization_dtype_map, skip_op_check, strip_debug_ops, weight_shard_size_bytes)
    141   if unsupported:
    142     raise ValueError('Unsupported Ops in the model before optimization\n' +
--> 143                      ', '.join(unsupported))
    144 
    145   # first pass of grappler optimization, this is needed for batch norm folding.

ValueError: Unsupported Ops in the model before optimization
StatefulPartitionedCall

[AIshutin]

Hi! I have the same error. @msahamed, did you manage to solve it or avoid somehow? (I need to run DistilBERT in a browser too)

[pyu10055]

@msahamed @Alshutin This fixed by this PR https://github.com/tensorflow/tfjs/pull/3685 You can wait for next release 2.1.0, we should be releasing this week or earlier next week.

[msahamed]

@msahamed @Alshutin This fixed by this PR #3685 You can wait for next release 2.1.0, we should be releasing this week or earlier next week.

Thank you for helping. It works with release 2.1.0.