1# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Keras model saving code."""
16
17from __future__ import absolute_import
18from __future__ import division
19from __future__ import print_function
20
21import six
22
23from tensorflow.python import tf2
24from tensorflow.python.keras.saving import hdf5_format
25from tensorflow.python.keras.saving import saving_utils
26from tensorflow.python.keras.saving.saved_model import load as saved_model_load
27from tensorflow.python.keras.saving.saved_model import load_context
28from tensorflow.python.keras.saving.saved_model import save as saved_model_save
29from tensorflow.python.keras.utils import generic_utils
30from tensorflow.python.keras.utils.io_utils import path_to_string
31from tensorflow.python.saved_model import loader_impl
32from tensorflow.python.util import keras_deps
33from tensorflow.python.util.tf_export import keras_export
34
35# pylint: disable=g-import-not-at-top
36try:
37  import h5py
38except ImportError:
39  h5py = None
40# pylint: enable=g-import-not-at-top
41
42
43@keras_export('keras.models.save_model')
44def save_model(model,
45               filepath,
46               overwrite=True,
47               include_optimizer=True,
48               save_format=None,
49               signatures=None,
50               options=None,
51               save_traces=True):
52  # pylint: disable=line-too-long
53  """Saves a model as a TensorFlow SavedModel or HDF5 file.
54
55  See the [Serialization and Saving guide](https://keras.io/guides/serialization_and_saving/)
56  for details.
57
58  Usage:
59
60  >>> model = tf.keras.Sequential([
61  ...     tf.keras.layers.Dense(5, input_shape=(3,)),
62  ...     tf.keras.layers.Softmax()])
63  >>> model.save('/tmp/model')
64  >>> loaded_model = tf.keras.models.load_model('/tmp/model')
65  >>> x = tf.random.uniform((10, 3))
66  >>> assert np.allclose(model.predict(x), loaded_model.predict(x))
67
68  The SavedModel and HDF5 file contains:
69
70  - the model's configuration (topology)
71  - the model's weights
72  - the model's optimizer's state (if any)
73
74  Thus models can be reinstantiated in the exact same state, without any of the
75  code used for model definition or training.
76
77  Note that the model weights may have different scoped names after being
78  loaded. Scoped names include the model/layer names, such as
79  `"dense_1/kernel:0"`. It is recommended that you use the layer properties to
80  access specific variables, e.g. `model.get_layer("dense_1").kernel`.
81
82  __SavedModel serialization format__
83
84  Keras SavedModel uses `tf.saved_model.save` to save the model and all
85  trackable objects attached to the model (e.g. layers and variables). The model
86  config, weights, and optimizer are saved in the SavedModel. Additionally, for
87  every Keras layer attached to the model, the SavedModel stores:
88
89    * the config and metadata -- e.g. name, dtype, trainable status
90    * traced call and loss functions, which are stored as TensorFlow subgraphs.
91
92  The traced functions allow the SavedModel format to save and load custom
93  layers without the original class definition.
94
95  You can choose to not save the traced functions by disabling the `save_traces`
96  option. This will decrease the time it takes to save the model and the
97  amount of disk space occupied by the output SavedModel. If you enable this
98  option, then you _must_ provide all custom class definitions when loading
99  the model. See the `custom_objects` argument in `tf.keras.models.load_model`.
100
101  Args:
102      model: Keras model instance to be saved.
103      filepath: One of the following:
104        - String or `pathlib.Path` object, path where to save the model
105        - `h5py.File` object where to save the model
106      overwrite: Whether we should overwrite any existing model at the target
107        location, or instead ask the user with a manual prompt.
108      include_optimizer: If True, save optimizer's state together.
109      save_format: Either 'tf' or 'h5', indicating whether to save the model
110        to Tensorflow SavedModel or HDF5. Defaults to 'tf' in TF 2.X, and 'h5'
111        in TF 1.X.
112      signatures: Signatures to save with the SavedModel. Applicable to the 'tf'
113        format only. Please see the `signatures` argument in
114        `tf.saved_model.save` for details.
115      options: (only applies to SavedModel format) `tf.saved_model.SaveOptions`
116        object that specifies options for saving to SavedModel.
117      save_traces: (only applies to SavedModel format) When enabled, the
118        SavedModel will store the function traces for each layer. This
119        can be disabled, so that only the configs of each layer are stored.
120        Defaults to `True`. Disabling this will decrease serialization time and
121        reduce file size, but it requires that all custom layers/models
122        implement a `get_config()` method.
123
124  Raises:
125      ImportError: If save format is hdf5, and h5py is not available.
126  """
127  # pylint: enable=line-too-long
128  from tensorflow.python.keras.engine import sequential  # pylint: disable=g-import-not-at-top
129
130  default_format = 'tf' if tf2.enabled() else 'h5'
131  save_format = save_format or default_format
132
133  filepath = path_to_string(filepath)
134
135  if (save_format == 'h5' or
136      (h5py is not None and isinstance(filepath, h5py.File)) or
137      saving_utils.is_hdf5_filepath(filepath)):
138    # TODO(b/130258301): add utility method for detecting model type.
139    if (not model._is_graph_network and  # pylint:disable=protected-access
140        not isinstance(model, sequential.Sequential)):
141      raise NotImplementedError(
142          'Saving the model to HDF5 format requires the model to be a '
143          'Functional model or a Sequential model. It does not work for '
144          'subclassed models, because such models are defined via the body of '
145          'a Python method, which isn\'t safely serializable. Consider saving '
146          'to the Tensorflow SavedModel format (by setting save_format="tf") '
147          'or using `save_weights`.')
148    hdf5_format.save_model_to_hdf5(
149        model, filepath, overwrite, include_optimizer)
150  else:
151    with generic_utils.SharedObjectSavingScope():
152      saved_model_save.save(model, filepath, overwrite, include_optimizer,
153                            signatures, options, save_traces)
154
155
156@keras_export('keras.models.load_model')
157def load_model(filepath, custom_objects=None, compile=True, options=None):  # pylint: disable=redefined-builtin
158  """Loads a model saved via `model.save()`.
159
160  Usage:
161
162  >>> model = tf.keras.Sequential([
163  ...     tf.keras.layers.Dense(5, input_shape=(3,)),
164  ...     tf.keras.layers.Softmax()])
165  >>> model.save('/tmp/model')
166  >>> loaded_model = tf.keras.models.load_model('/tmp/model')
167  >>> x = tf.random.uniform((10, 3))
168  >>> assert np.allclose(model.predict(x), loaded_model.predict(x))
169
170  Note that the model weights may have different scoped names after being
171  loaded. Scoped names include the model/layer names, such as
172  `"dense_1/kernel:0"`. It is recommended that you use the layer properties to
173  access specific variables, e.g. `model.get_layer("dense_1").kernel`.
174
175  Args:
176      filepath: One of the following:
177          - String or `pathlib.Path` object, path to the saved model
178          - `h5py.File` object from which to load the model
179      custom_objects: Optional dictionary mapping names
180          (strings) to custom classes or functions to be
181          considered during deserialization.
182      compile: Boolean, whether to compile the model
183          after loading.
184      options: Optional `tf.saved_model.LoadOptions` object that specifies
185        options for loading from SavedModel.
186
187  Returns:
188      A Keras model instance. If the original model was compiled, and saved with
189      the optimizer, then the returned model will be compiled. Otherwise, the
190      model will be left uncompiled. In the case that an uncompiled model is
191      returned, a warning is displayed if the `compile` argument is set to
192      `True`.
193
194  Raises:
195      ImportError: if loading from an hdf5 file and h5py is not available.
196      IOError: In case of an invalid savefile.
197  """
198  with generic_utils.SharedObjectLoadingScope():
199    with generic_utils.CustomObjectScope(custom_objects or {}):
200      with load_context.load_context(options):
201        if (h5py is not None and
202            (isinstance(filepath, h5py.File) or h5py.is_hdf5(filepath))):
203          return hdf5_format.load_model_from_hdf5(filepath, custom_objects,
204                                                  compile)
205
206        filepath = path_to_string(filepath)
207        if isinstance(filepath, six.string_types):
208          loader_impl.parse_saved_model(filepath)
209          return saved_model_load.load(filepath, compile, options)
210
211  raise IOError(
212      'Unable to load model. Filepath is not an hdf5 file (or h5py is not '
213      'available) or SavedModel.')
214
215# Inject the load_model function to keras_deps to remove the dependency
216# from TFLite to Keras.
217keras_deps.register_load_model_function(load_model)
218