无法加载 tensorflow（tf-agent）保存的模型

Question

I am creating a tf-agent DqnAgent in the following code:

tf_agent = dqn_agent.DqnAgent(
    train_env.time_step_spec(),
    train_env.action_spec(),
    q_network=q_net,
    optimizer=optimizer,
    td_errors_loss_fn=dqn_agent.element_wise_squared_loss,
    train_step_counter=train_step_counter

)

During the training loop I am saving this model with

tf.saved_model.save(tf_agent, saved_models_path)

Once trained, I want to load saved model with

if tf.saved_model.contains_saved_model(saved_models_path):
    tf_agent = tf.saved_model.load(saved_models_path)

This code will load the saved model only if the folder in saved_path contains one, the functions contains_saved_model(saved_models_path) returns True , so the model is loaded, but there is an excetion and the program crashes:

Traceback (most recent call last):
    File "/home/claudino/Projetos/dino-tf-agents/dino_ia/model/agent.py", line 50, in <module>
        tf_agent = tf.saved_model.load(saved_models_path)
    File "/home/claudino/Projetos/dino-tf-agents/venv/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 408, in load
        return load_internal(export_dir, tags)
    File "/home/claudino/Projetos/dino-tf-agents/venv/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 432, in load_internal
        export_dir)
    File "/home/claudino/Projetos/dino-tf-agents/venv/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 58, in __init__
        self._load_all()
    File "/home/claudino/Projetos/dino-tf-agents/venv/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 168, in _load_all
        slot_variable = optimizer_object.add_slot(
    AttributeError: '_UserObject' object has no attribute 'add_slot'

    Process finished with exit code 1

I navigated tensorflow code but couldn't find the problem. Anyone can help me?

I am using tf-agents-nightly because google's colaboratory source code don't work on tf-agents "stable" version (I am not sure tf-agents is really stable), and tryed the code with tensorflow 1.3 and 2.0.0-beta0 , same problem occurs.

Answer 1

Have you tried TensorFlow 2.7? That usually helps with this issue.

Something else that worked for me was to load the model this way (This assumes the model is a keras / tf.keras model):

try:
    model = tf.keras.models.load_model(model_dir)
except:
  load_options = tf.saved_model.LoadOptions(experimental_io_device= '/job:localhost')
  model = tf.saved_model.load(model_dir, options= load_options)

The try clause will lead to an exception, because load_model() requires a keras_metadata.pb file, which is absent when you save models with saved_model.save() .

However, running that clause will somehow make tf.saved_model.load() run without any problems. There's probably some sort of interaction happening in the background that I don't really understand, but it works for me and the ' no attribute add_slot ' error doesn't come up.