tf2.0 failed to save model if an input is used not used in the model directly

[henrysky]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): No
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Win 10 x64
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: No
• TensorFlow installed from (source or binary): Binary
• TensorFlow version (use command below): 2.0a0
• Python version: 3.7
• Bazel version (if compiling from source): N/A
• GCC/Compiler version (if compiling from source): N/A
• CUDA/cuDNN version: N/A
• GPU model and memory: N/A

Describe the current behavior

I have this use case where the model requires multiple inputs and some inputs go into the model layers and some inputs go into the loss function directly. It works fine in the past until I updated my package unit test to tf2.0 to test drive and find out tf2.0 failed test that run fine with tf1.x

I noticed that its mainly due to the tensor that goes directly to the loss function in this case is not presented in self._network_nodes in the line if node_key not in nn.keras_model._network_nodes: in tensorflow keras network.py.

Describe the expected behavior

I expect the model to be saved successfully with tf2.0 just as tf1.x

Code to reproduce the issue

import numpy as np

import tensorflow as tf
import tensorflow.keras as tfk
Sequence = tfk.utils.Sequence

Dense = tfk.layers.Dense
Input = tfk.layers.Input

Model = tfk.models.Model

def special_loss(weights):
    def special_loss_internal(true, pred):
        return (true - pred / weights)
    return special_loss_internal

# Model 1 which does not have Flatten
input_tensor1 = Input(shape=[200], name='input_1')
input_tensor2 = Input(shape=[10], name='input_2')
output_tensor1 = Dense(units=10, name='output_1')(input_tensor1)
output_tensor2 = Dense(units=10, name='output_2')(input_tensor1)

neuralnet = Model(inputs=[input_tensor1, input_tensor2], outputs=[output_tensor1, output_tensor2])
neuralnet.compile(loss=special_loss(input_tensor2), optimizer='adam')

neuralnet.save("test.h5")

Other info / logs

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
~\Anaconda3\envs\tf2\lib\site-packages\tensorflow\python\util\nest.py in pack_sequence_as(structure, flat_sequence, expand_composites)
    430     final_index, packed = _packed_nest_with_indices(structure, flat_sequence,
--> 431                                                     0, is_seq)
    432     if final_index < len(flat_sequence):

~\Anaconda3\envs\tf2\lib\site-packages\tensorflow\python\util\nest.py in _packed_nest_with_indices(structure, flat, index, is_seq)
    380     else:
--> 381       packed.append(flat[index])
    382       index += 1

IndexError: list index out of range

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
<ipython-input-1-6f88ca6a69df> in <module>
     26 neuralnet.compile(loss=special_loss(input_tensor2), optimizer='adam')
     27
---> 28 neuralnet.save("test.h5")

~\Anaconda3\envs\tf2\lib\site-packages\tensorflow\python\keras\engine\network.py in save(self, filepath, overwrite, include_optimizer)
   1312
   1313     from tensorflow.python.keras.models import save_model  # pylint: disable=g-import-not-at-top
-> 1314     save_model(self, filepath, overwrite, include_optimizer)
   1315
   1316   def save_weights(self, filepath, overwrite=True, save_format=None):

~\Anaconda3\envs\tf2\lib\site-packages\tensorflow\python\keras\saving\hdf5_format.py in save_model(model, filepath, overwrite, include_optimizer)
     99         {
    100             'class_name': model.__class__.__name__,
--> 101             'config': model.get_config()
    102         },
    103         default=serialization.get_json_type).encode('utf8')

~\Anaconda3\envs\tf2\lib\site-packages\tensorflow\python\keras\engine\network.py in get_config(self)
   1107       model_inputs.append(
   1108           tf_utils.ListWrapper([layer.name, new_node_index, tensor_index]))
-> 1109     model_inputs = nest.pack_sequence_as(self._nested_inputs, model_inputs)
   1110     model_inputs = tf_utils.convert_inner_node_data(model_inputs)
   1111     config['input_layers'] = model_inputs

~\Anaconda3\envs\tf2\lib\site-packages\tensorflow\python\util\nest.py in pack_sequence_as(structure, flat_sequence, expand_composites)
    438           "Could not pack sequence. Structure had %d elements, but "
    439           "flat_sequence had %d elements.  Structure: %s, flat_sequence: %s." %
--> 440           (len(flat_structure), len(flat_sequence), structure, flat_sequence))
    441   return _sequence_like(structure, packed)
    442

ValueError: Could not pack sequence. Structure had 2 elements, but flat_sequence had 1 elements.  Structure: [<tf.Tensor 'input_1:0' shape=(None, 200) dtype=float32>, <tf.Tensor 'input_2:0' shape=(None, 10) dtype=float32>], flat_sequence: [<tensorflow.python.keras.utils.tf_utils.ListWrapper object at 0x000001958CF2DDA0>].

[jvishnuvardhan]

I could reproduce the issue with TF2.0.0-alpha0 in Google colab. Thanks!

[henrysky]

is there a schedule for this issue to be fixed?? Because I really want this issue to be fixed before 2.0.0a1 so that I can see if there is any further issue in my CI tests.

[lixin9311]

Same problem with 1.14.RC1 on colab now.
It can't load model from an h5 file either.

[henrysky]

@jvishnuvardhan and @k-w-w:

Is there any ETA for a fix or any direction so I can contribute for a fix?? Thanks!

[jvishnuvardhan]

@henrysky As you are interested to contribute, You can raise a PR. Thanks!

[k-w-w]

Hi henrysky@ Thanks so much for implementing your fix! Sorry for not updating sooner! I submitted a change that addressed this issue earlier here: 401bbfc#diff-4ee308ea180d49ae81691348531a2b6d

[tensorflow-bot]

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

[TachikakaMin]

same issue in 1.14.0

[henrysky]

same issue in 1.14.0

It is just a one line patch, you can patch your 1.14.0 like me by yourself

401bbfc#diff-4ee308ea180d49ae81691348531a2b6d

[TachikakaMin]

same issue in 1.14.0

It is just a one line patch, you can patch your 1.14.0 like me by yourself

401bbfc#diff-4ee308ea180d49ae81691348531a2b6d

oh, thank you!
I have tried it but there will be another problem.

/usr/local/lib/python3.6/site-packages/tensorflow/python/keras/saving/hdf5_format.py in save_model_to_hdf5(model, filepath, overwrite, include_optimizer)
     99             'config': model.get_config()
    100         },
--> 101         default=serialization.get_json_type).encode('utf8')
    102 
    103     model_weights_group = f.create_group('model_weights')

/usr/local/lib/python3.6/json/__init__.py in dumps(obj, skipkeys, ensure_ascii, check_circular, allow_nan, cls, indent, separators, default, sort_keys, **kw)
    236         check_circular=check_circular, allow_nan=allow_nan, indent=indent,
    237         separators=separators, default=default, sort_keys=sort_keys,
--> 238         **kw).encode(obj)
    239 
    240 

/usr/local/lib/python3.6/json/encoder.py in encode(self, o)
    197         # exceptions aren't as detailed.  The list call should be roughly
    198         # equivalent to the PySequence_Fast that ''.join() would do.
--> 199         chunks = self.iterencode(o, _one_shot=True)
    200         if not isinstance(chunks, (list, tuple)):
    201             chunks = list(chunks)

/usr/local/lib/python3.6/json/encoder.py in iterencode(self, o, _one_shot)
    255                 self.key_separator, self.item_separator, self.sort_keys,
    256                 self.skipkeys, _one_shot)
--> 257         return _iterencode(o, 0)
    258 
    259 def _make_iterencode(markers, _default, _encoder, _indent, _floatstr,

/usr/local/lib/python3.6/site-packages/tensorflow/python/util/serialization.py in get_json_type(obj)
     67     return dict(obj)
     68 
---> 69   raise TypeError('Not JSON Serializable:', obj)

TypeError: ('Not JSON Serializable:', b'\n\x03Mul\x12\x03Mul\x1a\x0fconv2d7/Sigmoid\x1a\x0cdepth_result*\x07\n\x01T\x12\x020\x01')

And I use tf.saved_model.save to save the model and there is no error.

[k-w-w]

@TachikakaMin What model are you trying to save out (where is b'\n\x03Mul\x12\x03Mul\x1a\x0fconv2d7/Sigmoid\x1a\x0cdepth_result*\x07\n\x01T\x12\x020\x01' coming from?)

[OriAlpha]

same issue in 1.14.0

It is just a one line patch, you can patch your 1.14.0 like me by yourself

401bbfc#diff-4ee308ea180d49ae81691348531a2b6d

by this it gives error as

TypeError: ('Not JSON Serializable:', b'\n\rdense/Softmax\x12\x07Softmax\x1a\rdense/BiasAdd*\x07\n\x01T\x12\x020\x01')