Addons: Use tf.saved_model with tfa.activations.mish produce error: "Op type not registered 'Addons>Mish' in binary running"

[pikaliov]

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04
• TensorFlow version and how it was installed (source or binary): 2.1.0 (binary)
• TensorFlow-Addons version and how it was installed (source or binary): 0.8.3 (binary)
• Python version: 3.7
• Is GPU used? (yes/no): yes

Issue description:

I trained model using tfa.activations.mish for activation func of Dense layers.
For now I tried to load and serve trained model. However i get next error:

tensorflow.python.framework.errors_impl.NotFoundError: Op type not registered 'Addons>Mish' in binary running ...

Can you give me advices how to load this model?

Some usage inside layers:

def point_wise_feed_forward_network(d_model, dim_ffnn):
    """
    Point wise feed forward network which consists of two fully-connected layers with a ANN-activation in between.
    :param d_model:
    :param dim_ffnn:
    :return:
    """
    return tf.keras.Sequential([
      tf.keras.layers.Dense(dim_ffnn, activation=tfa.activations.mish),  # (batch_size, seq_len, dff)
      tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)
                                ])

Saving model:
```
...
self.model._set_inputs(inp_train, tf.convert_to_tensor(True))
....
self.model.save(path2model, save_format='tf')

Loading model:

loaded = tf.saved_model.load(path2model)
print(list(loaded.signatures.keys()))

Full error logs:

/usr/local/bin/python3.7m /home/pikaliov/Projects/transformerv3_bpe/src/tmp.py
2020-04-03 12:04:42.595320: W tensorflow/stream_executor/platform/default/dso_loader.cc:55] Could not load dynamic library 'libnvinfer.so.6'; dlerror: libnvinfer.so.6: cannot open shared object file: No such file or directory
2020-04-03 12:04:42.595372: W tensorflow/stream_executor/platform/default/dso_loader.cc:55] Could not load dynamic library 'libnvinfer_plugin.so.6'; dlerror: libnvinfer_plugin.so.6: cannot open shared object file: No such file or directory
2020-04-03 12:04:42.595378: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:30] Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
/usr/local/lib/python3.7/site-packages/pandas/compat/__init__.py:84: UserWarning: Could not import the lzma module. Your installed Python is incomplete. Attempting to use lzma compression will result in a RuntimeError.
warnings.warn(msg)
0.8.3
Traceback (most recent call last):
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 3708, in _get_op_def
return self._op_def_cache[type]
KeyError: 'Addons>Mish'

Traceback (most recent call last):
File "/home/pikaliov/Projects/transformerv3_bpe/src/tmp.py", line 6, in
loaded = tf.saved_model.load("/home/pikaliov/Projects/transformerv3_bpe/models_and_weights/transformer_asr_restore/model_dir/models")
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/saved_model/load.py", line 528, in load
return load_internal(export_dir, tags)
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/saved_model/load.py", line 552, in load_internal
export_dir)
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/saved_model/load.py", line 114, in __init__
meta_graph.graph_def.library))
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/saved_model/function_deserialization.py", line 312, in load_function_def_library
func_graph = function_def_lib.function_def_to_graph(copy)
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/function_def_to_graph.py", line 59, in function_def_to_graph
fdef, input_shapes)
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/function_def_to_graph.py", line 218, in function_def_to_graph_def
op_def = default_graph._get_op_def(node_def.op) # pylint: disable=protected-access
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 3712, in _get_op_def
c_api.TF_GraphGetOpDef(self._c_graph, compat.as_bytes(type), buf)
tensorflow.python.framework.errors_impl.NotFoundError: Op type not registered 'Addons>Mish' in binary running on pikaliov-MS-7C37. 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.

Process finished with exit code 1

```

[gabrieldemarmiesse]

Could you try tfa.register_all() before loading the model and report if it works?

EDIT: This won't actually work as register_all is only present in the master branch and hasn't been backported to the 0.8 branch. You are facing this bug actually: https://github.com/tensorflow/addons/issues/1151
If you use the nightly version of tensorflow addons: pip install tfa-nightly and do tfa.register_all() before loading the model, that should work.

[pikaliov]

Thank you, it works.

[ElliotVilhelm]

This seems to still be an issue when loading with the keras API, tf.keras.models.load_model. I worked around this error by saving only the weights of the model.