ValueError: Unsupported Ops in the model before optimization StatefulPartitionedCall

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

[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)

[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.