cannot use tf.keras.models.load_model to load SavedModel

[ericxsun]

Describe the bug
save the trained model in format of SavedModel and reload it using tf.keras.models.load_model won't suceed with the following exception:

TypeError: The two structures don't have the same nested structure.

First structure: type=tuple str=(({'combiner_output': TensorSpec(shape=(None, 128), dtype=tf.float32, name='input_1_1_combiner_output')}, {}), TensorSpec(shape=(None,), dtype=tf.int8, name='input_2'))

Second structure: type=tuple str=({'combiner_output': TensorSpec(shape=(None, 128), dtype=tf.float32, name='inputs/0/combiner_output')}, {})

More specifically: The two namedtuples don't have the same sequence type. First structure type=tuple str=({'combiner_output': TensorSpec(shape=(None, 128), dtype=tf.float32, name='input_1_1_combiner_output')}, {}) has type tuple, while second structure type=dict str={'combiner_output': TensorSpec(shape=(None, 128), dtype=tf.float32, name='inputs/0/combiner_output')} has type dict
Entire first structure:
(({'combiner_output': .}, {}), .)
Entire second structure:
({'combiner_output': .}, {})

To Reproduce
using the example mnist-simple-train

step 1: train the model with the following code

import logging
import shutil
import yaml

from ludwig.api import LudwigModel
from ludwig.datasets import mnist

# clean out prior results
shutil.rmtree('./results', ignore_errors=True)

# set up Python dictionary to hold model training parameters
with open('./config.yaml', 'r') as f:
    config = yaml.safe_load(f.read())

# Define Ludwig model object that drive model training
model = LudwigModel(config, logging_level=logging.INFO)

# load and split MNIST dataset
training_set, test_set, _ = mnist.load(split=True)

# initiate model training
(
    train_stats,  # training statistics
    _,
    output_directory  # location for training results saved to disk
) = model.train(
    training_set=training_set,
    test_set=test_set,
    experiment_name='simple_image_experiment',
    model_name='single_model',
    skip_save_processed_input=True
)

when finished, one could got the results:

Training: 100%|███████████████████████████████████████| 469/469 [00:23<00:00, 20.19it/s]
Evaluation train: 100%|███████████████████████████████| 469/469 [00:06<00:00, 67.49it/s]
Evaluation test : 100%|█████████████████████████████████| 79/79 [00:01<00:00, 64.11it/s]
Took 31.4292s
╒═════════╤════════╤════════════╤═════════════╕
│ label   │   loss │   accuracy │   hits_at_k │
╞═════════╪════════╪════════════╪═════════════╡
│ train   │ 0.1218 │     0.9697 │      0.9968 │
├─────────┼────────┼────────────┼─────────────┤
│ test    │ 0.1145 │     0.9709 │      0.9966 │
╘═════════╧════════╧════════════╧═════════════╛
╒════════════╤════════╕
│ combined   │   loss │
╞════════════╪════════╡
│ train      │ 0.1027 │
├────────────┼────────┤
│ test       │ 0.0926 │
╘════════════╧════════╛

step 2: save the model

>>> model.model.save_savedmodel('./savedmodel')

suceed with some warning, looks like:

WARNING:tensorflow:Skipping full serialization of Keras layer <ludwig.modules.reduction_modules.SequenceReducer object at 0x1396c34a8>, because it is not built.
WARNING:tensorflow:From /Users/qinluo/works/tools/env-py37/.venv/lib/python3.7/site-packages/tensorflow/python/training/tracking/tracking.py:111: Model.state_updates (from tensorflow.python.keras.engine.training) is deprecated and will be removed in a future version.
Instructions for updating:
This property should not be used in TensorFlow 2.0, as updates are applied automatically.
2021-01-17 22:58:40.627240: W tensorflow/python/util/util.cc:348] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
WARNING:tensorflow:From /Users/qinluo/works/tools/env-py37/.venv/lib/python3.7/site-packages/tensorflow/python/training/tracking/tracking.py:111: Layer.updates (from tensorflow.python.keras.engine.base_layer) is deprecated and will be removed in a future version.
Instructions for updating:
This property should not be used in TensorFlow 2.0, as updates are applied automatically.
WARNING:tensorflow:Skipping full serialization of Keras layer <ludwig.modules.reduction_modules.ReduceSum object at 0x1396c35c0>, because it is not built.
INFO:tensorflow:Assets written to: ./savedmodel/assets

step 3: reload the saved model

>>> import tensorflow as tf
>>> tf.keras.models.load_model('./savedmodel')

one could got the following error

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/works/tools/env-py37/.venv/lib/python3.7/site-packages/tensorflow/python/util/nest.py in assert_same_structure(nest1, nest2, check_types, expand_composites)
    394     _pywrap_utils.AssertSameStructure(nest1, nest2, check_types,
--> 395                                       expand_composites)
    396   except (ValueError, TypeError) as e:

TypeError: The two structures don't have the same nested structure.

First structure: type=tuple str=(({'combiner_output': TensorSpec(shape=(None, 128), dtype=tf.float32, name='input_1_1_combiner_output')}, {}), TensorSpec(shape=(None,), dtype=tf.int8, name='input_2'))

Second structure: type=tuple str=({'combiner_output': TensorSpec(shape=(None, 128), dtype=tf.float32, name='inputs/0/combiner_output')}, {})

More specifically: The two namedtuples don't have the same sequence type. First structure type=tuple str=({'combiner_output': TensorSpec(shape=(None, 128), dtype=tf.float32, name='input_1_1_combiner_output')}, {}) has type tuple, while second structure type=dict str={'combiner_output': TensorSpec(shape=(None, 128), dtype=tf.float32, name='inputs/0/combiner_output')} has type dict

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
<ipython-input-4-a155984253eb> in <module>
----> 1 tf.keras.models.load_model('./savedmodel')

~/works/tools/env-py37/.venv/lib/python3.7/site-packages/tensorflow/python/keras/saving/save.py in load_model(filepath, custom_objects, compile, options)
    185     if isinstance(filepath, six.string_types):
    186       loader_impl.parse_saved_model(filepath)
--> 187       return saved_model_load.load(filepath, compile, options)
    188
    189   raise IOError(

~/works/tools/env-py37/.venv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py in load(path, compile, options)
    119
    120   model = tf_load.load_internal(
--> 121       path, options=options, loader_cls=KerasObjectLoader)
    122
    123   # pylint: disable=protected-access

~/works/tools/env-py37/.venv/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py in load_internal(export_dir, tags, options, loader_cls)
    631       try:
    632         loader = loader_cls(object_graph_proto, saved_model_proto, export_dir,
--> 633                             ckpt_options)
    634       except errors.NotFoundError as err:
    635         raise FileNotFoundError(

~/works/tools/env-py37/.venv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py in __init__(self, *args, **kwargs)
    192     self._models_to_reconstruct = []
    193
--> 194     super(KerasObjectLoader, self).__init__(*args, **kwargs)
    195
    196     # Now that the node object has been fully loaded, and the checkpoint has

~/works/tools/env-py37/.venv/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py in __init__(self, object_graph_proto, saved_model_proto, export_dir, ckpt_options)
    128       self._concrete_functions[name] = _WrapperFunction(concrete_function)
    129
--> 130     self._load_all()
    131     self._restore_checkpoint()
    132

~/works/tools/env-py37/.venv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py in _load_all(self)
    219
    220     # Finish setting up layers and models. See function docstring for more info.
--> 221     self._finalize_objects()
    222
    223   @property

~/works/tools/env-py37/.venv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py in _finalize_objects(self)
    524         layers_revived_from_saved_model.append(node)
    525
--> 526     _finalize_saved_model_layers(layers_revived_from_saved_model)
    527     _finalize_config_layers(layers_revived_from_config)
    528

~/works/tools/env-py37/.venv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py in _finalize_saved_model_layers(layers)
    704         call_fn = _get_keras_attr(layer).call_and_return_conditional_losses
    705         if call_fn.input_signature is None:
--> 706           inputs = infer_inputs_from_restored_call_function(call_fn)
    707         else:
    708           inputs = call_fn.input_signature[0]

~/works/tools/env-py37/.venv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py in infer_inputs_from_restored_call_function(fn)
    983   for concrete in fn.concrete_functions[1:]:
    984     spec2 = concrete.structured_input_signature[0][0]
--> 985     spec = nest.map_structure(common_spec, spec, spec2)
    986   return spec
    987

~/works/tools/env-py37/.venv/lib/python3.7/site-packages/tensorflow/python/util/nest.py in map_structure(func, *structure, **kwargs)
    627   for other in structure[1:]:
    628     assert_same_structure(structure[0], other, check_types=check_types,
--> 629                           expand_composites=expand_composites)
    630
    631   flat_structure = [flatten(s, expand_composites) for s in structure]

~/works/tools/env-py37/.venv/lib/python3.7/site-packages/tensorflow/python/util/nest.py in assert_same_structure(nest1, nest2, check_types, expand_composites)
    400                   "Entire first structure:\n%s\n"
    401                   "Entire second structure:\n%s"
--> 402                   % (str(e), str1, str2))
    403
    404

TypeError: The two structures don't have the same nested structure.

First structure: type=tuple str=(({'combiner_output': TensorSpec(shape=(None, 128), dtype=tf.float32, name='input_1_1_combiner_output')}, {}), TensorSpec(shape=(None,), dtype=tf.int8, name='input_2'))

Second structure: type=tuple str=({'combiner_output': TensorSpec(shape=(None, 128), dtype=tf.float32, name='inputs/0/combiner_output')}, {})

More specifically: The two namedtuples don't have the same sequence type. First structure type=tuple str=({'combiner_output': TensorSpec(shape=(None, 128), dtype=tf.float32, name='input_1_1_combiner_output')}, {}) has type tuple, while second structure type=dict str={'combiner_output': TensorSpec(shape=(None, 128), dtype=tf.float32, name='inputs/0/combiner_output')} has type dict
Entire first structure:
(({'combiner_output': .}, {}), .)
Entire second structure:
({'combiner_output': .}, {})

Environment (please complete the following information):

• OS: Mac 10.13.6
• Python: python3.8 or python3.7 (both tested)
• Ludwig version: 0.3.2
• Tensorflow: 2.3.1

Additional context
maybe the line hidden = tf.reshape(hidden, [shape[0], -1]) in ludwig/encoders/image_encoders.py should be replaced by:

shape = hidden.shape
if shape[0] is not None:
    hidden = tf.reshape(hidden, [shape[0], -1])
else:
    hidden = tf.reshape(hidden, [-1, tf.math.reduce_prod(shape[1:])])

[jimthompson5802]

@ericxsun Thank you for reporting this issue.

It appears your issue is a variant of Issue #1070. A PR in underdevelopment to fix the issue. If you wish, you can monitor the status of PR #1074. Right now the PR addresses the issue you point out regarding tf.reshape() call. Additional work is need to address the warning messages you point out.

[ericxsun]

@ericxsun Thank you for reporting this issue.

It appears your issue is a variant of Issue #1070. A PR in underdevelopment to fix the issue. If you wish, you can monitor the status of PR #1074. Right now the PR addresses the issue you point out regarding tf.reshape() call. Additional work is need to address the warning messages you point out.

Thanks a lot.

PR #1074 solves the tf.reshape error, that's not my main point.

The problem is Cannot load saved model (in SavedModel) by only using tf.keras.models.load_model caused by TypeError: The two structures don't have the same nested structure.

Since I've fixed the shape TypeError locally: TypeError: Failed to convert object of type <class 'list'> to Tensor. Contents: [None, -1]. Consider casting elements to a supported type. as described in #1070

[jimthompson5802]

In Ludwig's unit test, we use tf.saved_model.load() to test restoring a saved model. Give that a try to see if resolves your issue.

[ericxsun]

In Ludwig's unit test, we use tf.saved_model.load() to test restoring a saved model. Give that a try to see if resolves your issue.

YES, you're right. Thanks a lot

tf.saved_model.load solved my problem.

[jimthompson5802]

@ericxsun Good to hear the last posting helped. If this addresses your concern, I'll close the issue.