tf.keras.models.load_model does not work for nsl.keras.AdversarialRegularization

[samedii]

Is there another intended way of saving model + optimizer?

2019-10-01 02:23:51.464271: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1304] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 7219 MB memory) -> physical GPU (device: 0, name: GeForce GTX 1070 with Max-Q Design, pci bus id: 0000:01:00.0, compute capability: 6.1)
Traceback (most recent call last):
  File "/home/richard/Documents/cloud/venv/.guild/runs/bca2e48d11674499949ff661421a60b2/.guild/sourcecode/retrain.py", line 52, in <module>
    model = load_model(config, pretrained_config)
  File "/home/richard/Documents/cloud/venv/.guild/runs/bca2e48d11674499949ff661421a60b2/.guild/sourcecode/cloud/model/load_model.py", line 32, in load_model
    model = tf.keras.models.load_model(directory + best_checkpoint_path)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/saving/save.py", line 150, in load_model
    return saved_model_load.load(filepath, compile)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/saving/saved_model/load.py", line 86, in load
    model = tf_load.load_internal(path, loader_cls=KerasObjectLoader)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/saved_model/load.py", line 541, in load_internal
    export_dir)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/saving/saved_model/load.py", line 103, in __init__
    self._finalize()
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/saving/saved_model/load.py", line 120, in _finalize
    inputs = infer_inputs_from_restored_call_function(call_fn)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/saving/saved_model/load.py", line 279, in infer_inputs_from_restored_call_function
    spec = nest.map_structure(common_spec, spec, spec2)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/util/nest.py", line 529, in map_structure
    expand_composites=expand_composites)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/util/nest.py", line 325, in assert_same_structure
    % (str(e), str1, str2))
ValueError: The two structures don't have the same nested structure.

First structure: type=TensorSpec str=TensorSpec(shape=(None, 256, 384, 3), dtype=tf.float32, name='feature')

Second structure: type=dict str={'feature': TensorSpec(shape=(None, 256, 384, 3), dtype=tf.float32, name='inputs/feature')}

More specifically: Substructure "type=dict str={'feature': TensorSpec(shape=(None, 256, 384, 3), dtype=tf.float32, name='inputs/feature')}" is a sequence, while substructure "type=TensorSpec str=TensorSpec(shape=(None, 256, 384, 3), dtype=tf.float32, name='feature')" is not
Entire first structure:
.
Entire second structure:
{'feature': .}

[samedii]

An alternative seems to be to use model.save_weights and:

1. Initialize model via code
2. Compile with the same optimizer
3. Use model.load_weights

Have not confirmed this works correctly but at least some optimizer weights are saved.

[csferng]

@samedii, would you mind sharing your goal of saving and restoring AdversarialRegularization models?

Given that AdversarialRegularization is mostly stateless (except the base model and compile options), a workaround could be to save the base model, create a new AdversarialRegularization from the restored base model, and compile with the same optimizer & loss.

[samedii]

Thanks for your answer! The only "need" is that It is normally a simple way of loading the model to continue training. It was an artifact from my code before adding AdversarialRegularization. Is your alternative better than the option I outlined above?

Sidenote: As you say, AdversarialRegularization should be stateless so I will try to recreate my issues with memory usage of AdversarialRegularization in a simple project as I find it strange that it more than doubles memory usage compared to when I'm not using it.

You can close this if you want, unless you want to use it to track documenting how to save/load a AdversarialRegularization model or something

[csferng]

After another look, I think save_weights+load_weight might be a better fit for continuing training, while saving the base model is an easier path for exporting an inference model.

Closing this now.