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"""Helper classes that list&validate all attributes to serialize to SavedModel."""
16
17import abc
18
19from tensorflow.python.keras.saving.saved_model import json_utils
20from tensorflow.python.keras.saving.saved_model import utils
21
22
23class SavedModelSaver(object, metaclass=abc.ABCMeta):
24 """Saver defining the methods and properties used to serialize Keras objects.
25 """
26
27 def __init__(self, obj):
28 self.obj = obj
29
30 @abc.abstractproperty
31 def object_identifier(self):
32 """String stored in object identifier field in the SavedModel proto.
33
34 Returns:
35 A string with the object identifier, which is used at load time.
36 """
37 raise NotImplementedError
38
39 @property
40 def tracking_metadata(self):
41 """String stored in metadata field in the SavedModel proto.
42
43 Returns:
44 A serialized JSON storing information necessary for recreating this layer.
45 """
46 # TODO(kathywu): check that serialized JSON can be loaded (e.g., if an
47 # object is in the python property)
48 return json_utils.Encoder().encode(self.python_properties)
49
50 def trackable_children(self, serialization_cache):
51 """Lists all Trackable children connected to this object."""
52 if not utils.should_save_traces():
53 return {}
54
55 children = self.objects_to_serialize(serialization_cache)
56 children.update(self.functions_to_serialize(serialization_cache))
57 return children
58
59 @abc.abstractproperty
60 def python_properties(self):
61 """Returns dictionary of python properties to save in the metadata.
62
63 This dictionary must be serializable and deserializable to/from JSON.
64
65 When loading, the items in this dict are used to initialize the object and
66 define attributes in the revived object.
67 """
68 raise NotImplementedError
69
70 @abc.abstractmethod
71 def objects_to_serialize(self, serialization_cache):
72 """Returns dictionary of extra checkpointable objects to serialize.
73
74 See `functions_to_serialize` for an explanation of this function's
75 effects.
76
77 Args:
78 serialization_cache: Dictionary passed to all objects in the same object
79 graph during serialization.
80
81 Returns:
82 A dictionary mapping attribute names to checkpointable objects.
83 """
84 raise NotImplementedError
85
86 @abc.abstractmethod
87 def functions_to_serialize(self, serialization_cache):
88 """Returns extra functions to include when serializing a Keras object.
89
90 Normally, when calling exporting an object to SavedModel, only the
91 functions and objects defined by the user are saved. For example:
92
93 ```
94 obj = tf.Module()
95 obj.v = tf.Variable(1.)
96
97 @tf.function
98 def foo(...): ...
99
100 obj.foo = foo
101
102 w = tf.Variable(1.)
103
104 tf.saved_model.save(obj, 'path/to/saved/model')
105 loaded = tf.saved_model.load('path/to/saved/model')
106
107 loaded.v # Variable with the same value as obj.v
108 loaded.foo # Equivalent to obj.foo
109 loaded.w # AttributeError
110 ```
111
112 Assigning trackable objects to attributes creates a graph, which is used for
113 both checkpointing and SavedModel serialization.
114
115 When the graph generated from attribute tracking is insufficient, extra
116 objects and functions may be added at serialization time. For example,
117 most models do not have their call function wrapped with a @tf.function
118 decorator. This results in `model.call` not being saved. Since Keras objects
119 should be revivable from the SavedModel format, the call function is added
120 as an extra function to serialize.
121
122 This function and `objects_to_serialize` is called multiple times when
123 exporting to SavedModel. Please use the cache to avoid generating new
124 functions and objects. A fresh cache is created for each SavedModel export.
125
126 Args:
127 serialization_cache: Dictionary passed to all objects in the same object
128 graph during serialization.
129
130 Returns:
131 A dictionary mapping attribute names to `Function` or
132 `ConcreteFunction`.
133 """
134 raise NotImplementedError