Can't load fine-tuned BERT model for inference

[wangler]

Have I written custom code : No
OS Platform and Distribution : macOs
TensorFlow version (use command below): tensorflow==2.4.0 and 2.5.0-dev20210121
Python version:Python 3.8 and Python 3.7

After successfully training and exporting a fine-tuned BERT classifier, I'm unable to re-load the model for inference in another session.

When I run
tf.saved_model.load("/path/to/saved_model")

or

with tf.device('/job:localhost'):
   reloaded_model = tf.saved_model.load("/path/to/saved_model")

I get:

  File "/Users/lucywang/Envs/dev/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 3936, in _get_op_def
    return self._op_def_cache[type]
KeyError: 'CaseFoldUTF8'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/lucywang/Envs/dev/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 895, in load_internal
    ckpt_options, filters)
  File "/Users/lucywang/Envs/dev/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 132, in __init__
    meta_graph.graph_def.library))
  File "/Users/lucywang/Envs/dev/lib/python3.7/site-packages/tensorflow/python/saved_model/function_deserialization.py", line 344, in load_function_def_library
    func_graph = function_def_lib.function_def_to_graph(copy)
  File "/Users/lucywang/Envs/dev/lib/python3.7/site-packages/tensorflow/python/framework/function_def_to_graph.py", line 64, in function_def_to_graph
    fdef, input_shapes)
  File "/Users/lucywang/Envs/dev/lib/python3.7/site-packages/tensorflow/python/framework/function_def_to_graph.py", line 228, in function_def_to_graph_def
    op_def = default_graph._get_op_def(node_def.op)  # pylint: disable=protected-access
  File "/Users/lucywang/Envs/dev/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 3941, in _get_op_def
    buf)
tensorflow.python.framework.errors_impl.NotFoundError: Op type not registered 'CaseFoldUTF8' in binary running on tw-mbp-lucywang. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
  File "/Users/lucywang/Envs/dev/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 864, in load
    return load_internal(export_dir, tags, options)["root"]
  File "/Users/lucywang/Envs/dev/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 898, in load_internal
    str(err) + "\n If trying to load on a different device from the "
FileNotFoundError: Op type not registered 'CaseFoldUTF8' in binary running on tw-mbp-lucywang. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed.
 If trying to load on a different device from the computational device, consider using setting the `experimental_io_device` option on tf.saved_model.LoadOptions to the io_device such as '/job:localhost'.

Importing tensorflow_text makes this issue go away (a solution suggested by #463), however this situation is not ideal as it'd require setting up a custom container with this tensorflow_text package at inference time. Is there a way to load a fine-tuned hub model for inference without having to import tensorflow_text?

I was able to reproduce this error when trying to load the model trained from this tutorial in an outside instance: https://www.tensorflow.org/tutorials/text/solve_glue_tasks_using_bert_on_tpu (colab). This tutorial explicitly packages the BERT preprocessor with the classifier model to make it self-sufficient for inference time, but even trying to load this exported model fails.

Any help would be much appreciated, thanks so much!

[arnoegw]

Hi @wangler, when it comes to text preprocessing at inference time, you have the choice between doing it inside your TensorFlow model – which requires TF.text ops to be installed – or outside your TensorFlow model – which requires equivalent logic to run on input text before it reaches the model, so some other part of your deployment will have to import code for that.

Both choices can make sense, depending on the circumstances. In terms of the tutorial, the latter choice comes down to exporting just classifier_model instead of the combined model_for_export, and then feeding it input_word_ids, input_mask and input_type_ids instead of text.

One advantage of bundling both pieces as shown in the tutorial is the encapsulation of preprocessing details by the exported model. The day you update your model to something with different preprocessing, it will be handy if you can just keep passing plain text.

All that said, my sense is the technical details are working as intended. In particular, the TensorFlow team has been very clear that simply adding useful ops to core TensorFlow for all sorts of problems is not feasible, which is why TF.text and other add-on libraries exist separately. As far as bug reports go, can we close this issue?

[wangler]

Got it, thanks for clarifying this!