Serialization problems in keras when using add_loss or '.tf' extension for saving

[nguerinjr]

System information

• WSL Win10 Ubuntu 18.04 (it also happens in a real ubuntu 18)
• tf-nightly-gpu-2.0-preview==2.0.0.dev20190802 (happens in cpu and gpu)
• python3.7

Describe the current behavior

I have a simple test to serialized and deserialize a model which has no compiled loss, just one added with add_loss.

Describe the expected behavior

Correct serialization and deserialization of the code

Code 1 to reproduce the issue

Occurs with both the pure loss function (which on earlier version would return json errors, but now seems ok), but also with the Lambda wrapper.

import tensorflow as tf
import tensorflow.keras as tf_k

model = tf_k.models.Sequential()
model.add(tf_k.layers.Conv2D(32, (3, 3), activation='relu', batch_size=32, input_shape=(28, 28, 1), padding='same', strides=2))
model.add(tf_k.layers.Conv2DTranspose(64, (3, 3), strides=2, padding='same'))
model.add(tf_k.layers.Conv2D(1, (3, 3), activation='sigmoid', padding='same'))
#loss = tf_k.layers.Lambda(lambda i: tf_k.losses.mse(*i))([model.inputs[0], model.outputs[0]])
loss = tf_k.losses.mean_squared_error(model.inputs[0], tf.zeros_like(model.inputs[0]))
model.add_loss(loss)
model.compile(optimizer='adam')
model.save('teste.h5')
model = tf.keras.models.load_model('teste.h5')

2019-08-03 03:38:05.624310: W tensorflow/stream_executor/platform/default/dso_loader.cc:55] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory
2019-08-03 03:38:05.624506: E tensorflow/stream_executor/cuda/cuda_driver.cc:318] failed call to cuInit: UNKNOWN ERROR (303)
2019-08-03 03:38:05.624652: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (DESKTOP-GUERINJR): /proc/driver/nvidia/version does not exist
2019-08-03 03:38:05.624991: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2019-08-03 03:38:05.634458: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 3408000000 Hz
2019-08-03 03:38:05.636116: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x7fffe06c0120 executing computations on platform Host. Devices:
2019-08-03 03:38:05.636259: I tensorflow/compiler/xla/service/service.cc:175]   StreamExecutor device (0): Host, Default Version
WARNING: Logging before flag parsing goes to stderr.
W0803 03:38:05.683956 139979377280832 training_utils.py:1320] Output conv2d_1 missing from loss dictionary. We assume this was done on purpose. The fit and evaluate APIs will not be expecting any data to be passed to conv2d_1.
Traceback (most recent call last):
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 1554, in _create_c_op
    c_op = c_api.TF_FinishOperation(op_desc)
tensorflow.python.framework.errors_impl.InvalidArgumentError: 1 inputs specified of 2 inputs in Op while building NodeDef 'tf_op_layer_SquaredDifference/SquaredDifference' using Op<name=SquaredDifference; signature=x:T, y:T -> z:T; attr=T:type,allowed=[DT_BFLOAT16, DT_HALF, DT_FLOAT, DT_DOUBLE, DT_INT32, DT_INT64, DT_COMPLEX64, DT_COMPLEX128]; is_commutative=true>

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/nguerinjr/Documents/deep_coding_project/img_common/teste.py", line 14, in <module>
    model = tf.keras.models.load_model('teste.h5')
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/save.py", line 138, in load_model
    return hdf5_format.load_model_from_hdf5(filepath, custom_objects, compile)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/hdf5_format.py", line 162, in load_model_from_hdf5
    custom_objects=custom_objects)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/model_config.py", line 55, in model_from_config
    return deserialize(config, custom_objects=custom_objects)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/layers/serialization.py", line 98, in deserialize
    printable_module_name='layer')
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/utils/generic_utils.py", line 191, in deserialize_keras_object
    list(custom_objects.items())))
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/sequential.py", line 369, in from_config
    model.add(layer)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/training/tracking/base.py", line 457, in _method_wrapper
    result = method(self, *args, **kwargs)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/sequential.py", line 195, in add
    output_tensor = layer(self.outputs[0])
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 799, in __call__
    outputs = call_fn(cast_inputs, *args, **kwargs)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 2507, in call
    return self._make_op(inputs)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 2530, in _make_op
    c_op = ops._create_c_op(graph, node_def, inputs, control_inputs=[])
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 1557, in _create_c_op
    raise ValueError(str(e))
ValueError: 1 inputs specified of 2 inputs in Op while building NodeDef 'tf_op_layer_SquaredDifference/SquaredDifference' using Op<name=SquaredDifference; signature=x:T, y:T -> z:T; attr=T:type,allowed=[DT_BFLOAT16, DT_HALF, DT_FLOAT, DT_DOUBLE, DT_INT32, DT_INT64, DT_COMPLEX64, DT_COMPLEX128]; is_commutative=true>

Process finished with exit code 1

Code 2 to reproduce the issue

Putting the '.tf' suffix, only changes the error message. Here's an example

import tensorflow as tf
import tensorflow.keras as tf_k

model = tf_k.models.Sequential()
model.add(tf_k.layers.Conv2D(32, (3, 3), activation='relu', batch_size=32, input_shape=(28, 28, 1), padding='same', strides=2))
model.add(tf_k.layers.Conv2DTranspose(64, (3, 3), strides=2, padding='same'))
model.add(tf_k.layers.Conv2D(1, (3, 3), activation='sigmoid', padding='same'))

#loss = tf_k.layers.Lambda(lambda i: tf_k.losses.mse(*i))([model.inputs[0], model.outputs[0]])
loss = tf_k.losses.mean_squared_error(model.inputs[0], tf.zeros_like(model.inputs[0]))
model.add_loss(loss)
model.compile(optimizer='adam')
model.save('teste.tf')
model = tf.keras.models.load_model('teste.tf')

2019-08-03 03:38:42.774544: W tensorflow/stream_executor/platform/default/dso_loader.cc:55] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory
2019-08-03 03:38:42.774737: E tensorflow/stream_executor/cuda/cuda_driver.cc:318] failed call to cuInit: UNKNOWN ERROR (303)
2019-08-03 03:38:42.774932: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (DESKTOP-GUERINJR): /proc/driver/nvidia/version does not exist
2019-08-03 03:38:42.775273: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2019-08-03 03:38:42.784443: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 3408000000 Hz
2019-08-03 03:38:42.786142: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x7fffc04e4040 executing computations on platform Host. Devices:
2019-08-03 03:38:42.786289: I tensorflow/compiler/xla/service/service.cc:175]   StreamExecutor device (0): Host, Default Version
WARNING: Logging before flag parsing goes to stderr.
W0803 03:38:42.835703 140157751789376 training_utils.py:1320] Output conv2d_1 missing from loss dictionary. We assume this was done on purpose. The fit and evaluate APIs will not be expecting any data to be passed to conv2d_1.
2019-08-03 03:38:43.281262: W tensorflow/python/util/util.cc:288] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
W0803 03:38:43.369712 140157751789376 deprecation.py:506] From /home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1784: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.
Traceback (most recent call last):
  File "/home/nguerinjr/Documents/deep_coding_project/img_common/teste.py", line 14, in <module>
    model = tf.keras.models.load_model('teste.tf')
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/save.py", line 142, in load_model
    return saved_model_load.load(filepath, compile)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/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/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/saved_model/load.py", line 541, in load_internal
    export_dir)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/saved_model/load.py", line 103, in __init__
    self._finalize()
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/saved_model/load.py", line 127, in _finalize
    node.add(layer)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/training/tracking/base.py", line 457, in _method_wrapper
    result = method(self, *args, **kwargs)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/sequential.py", line 195, in add
    output_tensor = layer(self.outputs[0])
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 799, in __call__
    outputs = call_fn(cast_inputs, *args, **kwargs)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/saved_model/utils.py", line 57, in return_outputs_and_add_losses
    outputs, losses = fn(inputs, *args, **kwargs)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py", line 439, in __call__
    self._initialize(args, kwds, add_initializers_to=initializer_map)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py", line 382, in _initialize
    *args, **kwds))
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py", line 1806, in _get_concrete_function_internal_garbage_collected
    graph_function, _, _ = self._maybe_define_function(args, kwargs)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py", line 2106, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py", line 1997, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/framework/func_graph.py", line 884, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py", line 325, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/saved_model/function_deserialization.py", line 262, in restored_function_body
    "\n\n".join(signature_descriptions)))
ValueError: Could not find matching function to call loaded from the SavedModel. Got:
  Positional arguments (1 total):
    * Tensor("inputs:0", shape=(32, 28, 28, 1), dtype=float32)
  Keyword arguments: {}

Expected these arguments to match one of the following 1 option(s):

Option 1:
  Positional arguments (1 total):
    * [TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32, name='inputs/0')]
  Keyword arguments: {}

Process finished with exit code 1

Is there something wrong with this simple example? Something I'm missing. Or it's just that the add_loss still have some problems wrt serialization / deserialization?

I perceived a bunch of errors have been fixed since my last report about this custom uses of keras (#30378). For example, I know a workaround for this, which seems ok in this nightly version, is to use a custom compiled loss, which was not working at that time, as another issue I had reported (#30384)

import tensorflow as tf
import tensorflow.keras as tf_k

model = tf_k.models.Sequential()
model.add(tf_k.layers.Conv2D(32, (3, 3), activation='relu', batch_size=32, input_shape=(28, 28, 1), padding='same', strides=2))
model.add(tf_k.layers.Conv2DTranspose(64, (3, 3), strides=2, padding='same'))
model.add(tf_k.layers.Conv2D(1, (3, 3), activation='sigmoid', padding='same'))

def my_loss(y_true, y_pred):
    return y_true

def my_loss2(y_true, y_pred):
    return y_pred

model.compile(optimizer='adam', loss=[my_loss])
model.save('teste.h5')

model = tf.keras.models.load_model('teste.h5', custom_objects={'my_loss': my_loss, 'my_loss2': my_loss2})
print(model)

/home/nguerinjr/Documents/deep_coding_project/venv/bin/python /home/nguerinjr/Documents/deep_coding_project/img_common/teste.py
2019-08-03 03:46:57.744222: W tensorflow/stream_executor/platform/default/dso_loader.cc:55] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory
2019-08-03 03:46:57.744406: E tensorflow/stream_executor/cuda/cuda_driver.cc:318] failed call to cuInit: UNKNOWN ERROR (303)
2019-08-03 03:46:57.744557: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (DESKTOP-GUERINJR): /proc/driver/nvidia/version does not exist
2019-08-03 03:46:57.744891: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2019-08-03 03:46:57.755141: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 3408000000 Hz
2019-08-03 03:46:57.756825: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x7fffea867eb0 executing computations on platform Host. Devices:
2019-08-03 03:46:57.757004: I tensorflow/compiler/xla/service/service.cc:175]   StreamExecutor device (0): Host, Default Version
<tensorflow.python.keras.engine.sequential.Sequential object at 0x7f4e203ee940>

Process finished with exit code 0

But, as an extension to these tests, I've noticed the '.tf' version does not work correctly, which seems another related bug:

import tensorflow as tf
import tensorflow.keras as tf_k

model = tf_k.models.Sequential()
model.add(tf_k.layers.Conv2D(32, (3, 3), activation='relu', batch_size=32, input_shape=(28, 28, 1), padding='same', strides=2))
model.add(tf_k.layers.Conv2DTranspose(64, (3, 3), strides=2, padding='same'))
model.add(tf_k.layers.Conv2D(1, (3, 3), activation='sigmoid', padding='same'))

def my_loss(y_true, y_pred):
    return y_true

def my_loss2(y_true, y_pred):
    return y_pred

model.compile(optimizer='adam', loss=[my_loss])
model.save('teste.tf')

model = tf.keras.models.load_model('teste.tf', custom_objects={'my_loss': my_loss, 'my_loss2': my_loss2})
print(model)

/home/nguerinjr/Documents/deep_coding_project/venv/bin/python /home/nguerinjr/Documents/deep_coding_project/img_common/teste.py
2019-08-03 03:47:16.464514: W tensorflow/stream_executor/platform/default/dso_loader.cc:55] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory
2019-08-03 03:47:16.464702: E tensorflow/stream_executor/cuda/cuda_driver.cc:318] failed call to cuInit: UNKNOWN ERROR (303)
2019-08-03 03:47:16.464859: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (DESKTOP-GUERINJR): /proc/driver/nvidia/version does not exist
2019-08-03 03:47:16.465200: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2019-08-03 03:47:16.474429: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 3408000000 Hz
2019-08-03 03:47:16.476021: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x7fffbb87a7b0 executing computations on platform Host. Devices:
2019-08-03 03:47:16.476185: I tensorflow/compiler/xla/service/service.cc:175]   StreamExecutor device (0): Host, Default Version
2019-08-03 03:47:16.873280: W tensorflow/python/util/util.cc:288] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
WARNING: Logging before flag parsing goes to stderr.
W0803 03:47:16.921350 140320731367232 deprecation.py:506] From /home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1784: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.
Traceback (most recent call last):
  File "/home/nguerinjr/Documents/deep_coding_project/img_common/teste.py", line 25, in <module>
    model = tf.keras.models.load_model('teste.tf', custom_objects={'my_loss': my_loss, 'my_loss2': my_loss2})
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/save.py", line 142, in load_model
    return saved_model_load.load(filepath, compile)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/saved_model/load.py", line 93, in load
    model._training_config))  # pylint: disable=protected-access
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/training/tracking/base.py", line 457, in _method_wrapper
    result = method(self, *args, **kwargs)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 340, in compile
    self.loss, self.output_names)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_utils.py", line 1329, in prepare_loss_functions
    loss_functions = nest.map_structure(get_loss_function, loss)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/util/nest.py", line 523, in map_structure
    structure[0], [func(*x) for x in entries],
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/util/nest.py", line 523, in <listcomp>
    structure[0], [func(*x) for x in entries],
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_utils.py", line 1086, in get_loss_function
    loss_fn = losses.get(loss)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/losses.py", line 1166, in get
    return deserialize(identifier)
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/losses.py", line 1157, in deserialize
    printable_module_name='loss function')
  File "/home/nguerinjr/Documents/deep_coding_project/venv/lib/python3.7/site-packages/tensorflow_core/python/keras/utils/generic_utils.py", line 210, in deserialize_keras_object
    raise ValueError('Unknown ' + printable_module_name + ':' + object_name)
ValueError: Unknown loss function:my_loss

Process finished with exit code 1

The real point about my interest in the add_loss it that's the only way, at least I think, to use losses in a flexible way. I still need very flexible losses. add_loss seems an interesting way to do this.

[ravikyram]

@nguerinjr
Can you please check with the nightly version pip install tf-nightly-2.0-preview and see if the problem still persists.
Thanks!

[nguerinjr]

@ravikyram

Yes. All problems occur exactly the same way in version installed with pip3 install tf-nightly-2.0-preview which is, today, 2.0.0.dev20190808

[ravikyram]

I have tried on colab with TF version nightly versions 2.0.0.dev20190802 , 2.0.0.dev20190808 and was able to reproduce the issue.Please, find the gist here. Thanks!

[amahendrakar]

@nguerinjr,
Two of the code samples which you have provided work without any issues with the latest TF-nightly i.e. TF v2.2.0-dev20200416.

For the first two samples, I'm facing an error stating IndexError: list index out of range.

Please find the gist of it here. Thanks!

[google-ml-butler]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you.

[amahendrakar]

@nguerinjr,
Any updates regarding this issue? Thanks!

[amahendrakar]

Automatically closing due to lack of recent activity. Please update the issue when new information becomes available, and we will reopen the issue. Thanks!

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No