Can't load Keras model using RectifiedAdam optimizer

Viewed 1122

Trying to load a Keras model using tf.keras.models.load_model I get the following error:

import tensorflow as tf
from tensorflow_addons.optimizers import RectifiedAdam
model = tf.keras.models.load_model('model', custom_objects = {'RectifiedAdam' : RectifiedAdam})

Error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/mylib/python3.7/site-packages/tensorflow_core/python/keras/saving/save.py", line 150, in load_model
    return saved_model_load.load(filepath, compile)
  File "/mylib/python3.7/site-packages/tensorflow_core/python/keras/saving/saved_model/load.py", line 99, in load
    training_config))
  File "/mylib/python3.7/site-packages/tensorflow_core/python/keras/saving/saving_utils.py", line 229, in compile_args_from_training_config
    optimizer_config, custom_objects=custom_objects)
  File "/mylib/python3.7/site-packages/tensorflow_core/python/keras/optimizers.py", line 819, in deserialize
    printable_module_name='optimizer')
  File "/mylib/python3.7/site-packages/tensorflow_core/python/keras/utils/generic_utils.py", line 292, in deserialize_keras_object
    config, module_objects, custom_objects, printable_module_name)
  File "/mylib/python3.7/site-packages/tensorflow_core/python/keras/utils/generic_utils.py", line 250, in class_and_config_for_serialized_keras_object
    raise ValueError('Unknown ' + printable_module_name + ': ' + class_name)
ValueError: Unknown optimizer: RectifiedAdam

I can load the model with compile to False tf.keras.models.load_model('model', compile=False) and then compile it again with RectifiedAdam optimizer (as suggested here: https://stackoverflow.com/a/56565801) - however that is not ideal...

So any ideas on, what I'm doing wrong?

2 Answers

One quick hack around this is to manually assign RectifiedAdam to an object in the Tensorflow namespace:

import tensorflow as tf
from tensorflow_addons.optimizers import RectifiedAdam

tf.keras.optimizers.RectifiedAdam = RectifiedAdam
...

or do something like this:

models.load_model('myModel.h5', custom_objects={'MyOptimizer': MyOptimizer})
Related