1# Copyright 2017 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"""SavedModel simple save functionality."""
16
17from __future__ import absolute_import
18from __future__ import division
19from __future__ import print_function
20
21from tensorflow.python.framework import ops
22from tensorflow.python.saved_model import builder
23from tensorflow.python.saved_model import signature_constants
24from tensorflow.python.saved_model import signature_def_utils
25from tensorflow.python.saved_model import tag_constants
26from tensorflow.python.util import deprecation
27from tensorflow.python.util.tf_export import tf_export
28
29
30@tf_export(v1=['saved_model.simple_save'])
31@deprecation.deprecated(
32    None,
33    'This function will only be available through the v1 compatibility '
34    'library as tf.compat.v1.saved_model.simple_save.')
35def simple_save(session, export_dir, inputs, outputs, legacy_init_op=None):
36  """Convenience function to build a SavedModel suitable for serving.
37
38  In many common cases, saving models for serving will be as simple as:
39
40      simple_save(session,
41                  export_dir,
42                  inputs={"x": x, "y": y},
43                  outputs={"z": z})
44
45  Although in many cases it's not necessary to understand all of the many ways
46      to configure a SavedModel, this method has a few practical implications:
47    - It will be treated as a graph for inference / serving (i.e. uses the tag
48      `tag_constants.SERVING`)
49    - The SavedModel will load in TensorFlow Serving and supports the
50      [Predict
51      API](https://github.com/tensorflow/serving/blob/master/tensorflow_serving/apis/predict.proto).
52      To use the Classify, Regress, or MultiInference APIs, please
53      use either
54      [tf.Estimator](https://www.tensorflow.org/api_docs/python/tf/estimator/Estimator)
55      or the lower level
56      [SavedModel
57      APIs](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/README.md).
58    - Some TensorFlow ops depend on information on disk or other information
59      called "assets". These are generally handled automatically by adding the
60      assets to the `GraphKeys.ASSET_FILEPATHS` collection. Only assets in that
61      collection are exported; if you need more custom behavior, you'll need to
62      use the
63      [SavedModelBuilder](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/builder.py).
64
65  More information about SavedModel and signatures can be found here:
66  https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/README.md.
67
68  Args:
69    session: The TensorFlow session from which to save the meta graph and
70        variables.
71    export_dir: The path to which the SavedModel will be stored.
72    inputs: dict mapping string input names to tensors. These are added
73        to the SignatureDef as the inputs.
74    outputs:  dict mapping string output names to tensors. These are added
75        to the SignatureDef as the outputs.
76    legacy_init_op: Legacy support for op or group of ops to execute after the
77        restore op upon a load.
78  """
79  signature_def_map = {
80      signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
81          signature_def_utils.predict_signature_def(inputs, outputs)
82  }
83  b = builder.SavedModelBuilder(export_dir)
84  b.add_meta_graph_and_variables(
85      session,
86      tags=[tag_constants.SERVING],
87      signature_def_map=signature_def_map,
88      assets_collection=ops.get_collection(ops.GraphKeys.ASSET_FILEPATHS),
89      main_op=legacy_init_op,
90      clear_devices=True)
91  b.save()
92