Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/keras/src/export/export_lib.py: 17%
161 statements
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-03 07:57 +0000
1# Copyright 2023 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"""Library for exporting inference-only Keras models/layers."""
17import tensorflow.compat.v2 as tf
18from tensorflow.python.util.tf_export import keras_export
20from keras.src.engine import base_layer
21from keras.src.engine import functional
22from keras.src.engine import sequential
23from keras.src.utils import io_utils
26@keras_export("keras.export.ExportArchive")
27class ExportArchive(tf.__internal__.tracking.AutoTrackable):
28 """ExportArchive is used to write SavedModel artifacts (e.g. for inference).
30 If you have a Keras model or layer that you want to export as SavedModel for
31 serving (e.g. via TensorFlow-Serving), you can use `ExportArchive`
32 to configure the different serving endpoints you need to make available,
33 as well as their signatures. Simply instantiate an `ExportArchive`,
34 use `track()` to register the layer(s) or model(s) to be used,
35 then use the `add_endpoint()` method to register a new serving endpoint.
36 When done, use the `write_out()` method to save the artifact.
38 The resulting artifact is a SavedModel and can be reloaded via
39 `tf.saved_model.load`.
41 Examples:
43 Here's how to export a model for inference.
45 ```python
46 export_archive = ExportArchive()
47 export_archive.track(model)
48 export_archive.add_endpoint(
49 name="serve",
50 fn=model.call,
51 input_signature=[tf.TensorSpec(shape=(None, 3), dtype=tf.float32)],
52 )
53 export_archive.write_out("path/to/location")
55 # Elsewhere, we can reload the artifact and serve it.
56 # The endpoint we added is available as a method:
57 serving_model = tf.saved_model.load("path/to/location")
58 outputs = serving_model.serve(inputs)
59 ```
61 Here's how to export a model with one endpoint for inference and one
62 endpoint for a training-mode forward pass (e.g. with dropout on).
64 ```python
65 export_archive = ExportArchive()
66 export_archive.track(model)
67 export_archive.add_endpoint(
68 name="call_inference",
69 fn=lambda x: model.call(x, training=False),
70 input_signature=[tf.TensorSpec(shape=(None, 3), dtype=tf.float32)],
71 )
72 export_archive.add_endpoint(
73 name="call_training",
74 fn=lambda x: model.call(x, training=True),
75 input_signature=[tf.TensorSpec(shape=(None, 3), dtype=tf.float32)],
76 )
77 export_archive.write_out("path/to/location")
78 ```
80 **Note on resource tracking:**
82 `ExportArchive` is able to automatically track all `tf.Variables` used
83 by its endpoints, so most of the time calling `.track(model)`
84 is not strictly required. However, if your model uses lookup layers such
85 as `IntegerLookup`, `StringLookup`, or `TextVectorization`,
86 it will need to be tracked explicitly via `.track(model)`.
88 Explicit tracking is also required if you need to be able to access
89 the properties `variables`, `trainable_variables`, or
90 `non_trainable_variables` on the revived archive.
91 """
93 def __init__(self):
94 self._endpoint_names = []
95 self._endpoint_signatures = {}
96 self.tensorflow_version = tf.__version__
97 self.variables = []
98 self.trainable_variables = []
99 self.non_trainable_variables = []
101 @tf.__internal__.tracking.no_automatic_dependency_tracking
102 def track(self, layer):
103 """Track the variables (and other resources) of a layer or model."""
104 if not isinstance(layer, base_layer.Layer):
105 raise ValueError(
106 "Invalid layer type. Expected an instance of "
107 "`keras.layers.Layer` or `keras.Model`. "
108 f"Received instead an object of type '{type(layer)}'. "
109 f"Object received: {layer}"
110 )
111 if not layer.built:
112 raise ValueError(
113 "The layer provided has not yet been built. "
114 "It must be built before export."
115 )
117 # Layers in `_tracked` are not part of the trackables that get saved,
118 # because we're creating the attribute in a
119 # no_automatic_dependency_tracking scope.
120 if not hasattr(self, "_tracked"):
121 self._tracked = []
122 self._tracked.append(layer)
124 # Variables in the lists below are actually part of the trackables
125 # that get saved, because the lists are created in __init__.
126 self.variables += layer.variables
127 self.trainable_variables += layer.trainable_variables
128 self.non_trainable_variables += layer.non_trainable_variables
130 def add_endpoint(self, name, fn, input_signature=None):
131 """Register a new serving endpoint.
133 Arguments:
134 name: Str, name of the endpoint.
135 fn: A function. It should only leverage resources
136 (e.g. `tf.Variable` objects or `tf.lookup.StaticHashTable`
137 objects) that are available on the models/layers
138 tracked by the `ExportArchive` (you can call `.track(model)`
139 to track a new model).
140 The shape and dtype of the inputs to the function must be
141 known. For that purpose, you can either 1) make sure that
142 `fn` is a `tf.function` that has been called at least once, or
143 2) provide an `input_signature` argument that specifies the
144 shape and dtype of the inputs (see below).
145 input_signature: Used to specify the shape and dtype of the
146 inputs to `fn`. List of `tf.TensorSpec` objects (one
147 per positional input argument of `fn`). Nested arguments are
148 allowed (see below for an example showing a Functional model
149 with 2 input arguments).
151 Example:
153 Adding an endpoint using the `input_signature` argument when the
154 model has a single input argument:
156 ```python
157 export_archive = ExportArchive()
158 export_archive.track(model)
159 export_archive.add_endpoint(
160 name="serve",
161 fn=model.call,
162 input_signature=[tf.TensorSpec(shape=(None, 3), dtype=tf.float32)],
163 )
164 ```
166 Adding an endpoint using the `input_signature` argument when the
167 model has two positional input arguments:
169 ```python
170 export_archive = ExportArchive()
171 export_archive.track(model)
172 export_archive.add_endpoint(
173 name="serve",
174 fn=model.call,
175 input_signature=[
176 tf.TensorSpec(shape=(None, 3), dtype=tf.float32),
177 tf.TensorSpec(shape=(None, 4), dtype=tf.float32),
178 ],
179 )
180 ```
182 Adding an endpoint using the `input_signature` argument when the
183 model has one input argument that is a list of 2 tensors (e.g.
184 a Functional model with 2 inputs):
186 ```python
187 model = keras.Model(inputs=[x1, x2], outputs=outputs)
189 export_archive = ExportArchive()
190 export_archive.track(model)
191 export_archive.add_endpoint(
192 name="serve",
193 fn=model.call,
194 input_signature=[
195 [
196 tf.TensorSpec(shape=(None, 3), dtype=tf.float32),
197 tf.TensorSpec(shape=(None, 4), dtype=tf.float32),
198 ],
199 ],
200 )
201 ```
203 This also works with dictionary inputs:
205 ```python
206 model = keras.Model(inputs={"x1": x1, "x2": x2}, outputs=outputs)
208 export_archive = ExportArchive()
209 export_archive.track(model)
210 export_archive.add_endpoint(
211 name="serve",
212 fn=model.call,
213 input_signature=[
214 {
215 "x1": tf.TensorSpec(shape=(None, 3), dtype=tf.float32),
216 "x2": tf.TensorSpec(shape=(None, 4), dtype=tf.float32),
217 },
218 ],
219 )
220 ```
222 Adding an endpoint that is a `tf.function`:
224 ```python
225 @tf.function()
226 def serving_fn(x):
227 return model(x)
229 # The function must be traced, i.e. it must be called at least once.
230 serving_fn(tf.random.normal(shape=(2, 3)))
232 export_archive = ExportArchive()
233 export_archive.track(model)
234 export_archive.add_endpoint(name="serve", fn=serving_fn)
235 ```
236 """
237 if name in self._endpoint_names:
238 raise ValueError(f"Endpoint name '{name}' is already taken.")
240 if input_signature:
241 decorated_fn = tf.function(fn, input_signature=input_signature)
242 self._endpoint_signatures[name] = input_signature
243 else:
244 if isinstance(fn, tf.types.experimental.GenericFunction):
245 if not fn._list_all_concrete_functions():
246 raise ValueError(
247 f"The provided tf.function '{fn}' "
248 "has never been called. "
249 "To specify the expected shape and dtype "
250 "of the function's arguments, "
251 "you must either provide a function that "
252 "has been called at least once, or alternatively pass "
253 "an `input_signature` argument in `add_endpoint()`."
254 )
255 decorated_fn = fn
256 else:
257 raise ValueError(
258 "If the `fn` argument provided is not a `tf.function`, "
259 "you must provide an `input_signature` argument to "
260 "specify the shape and dtype of the function arguments. "
261 "Example:\n\n"
262 "export_archive.add_endpoint(\n"
263 " name='call',\n"
264 " fn=model.call,\n"
265 " input_signature=[\n"
266 " tf.TensorSpec(\n"
267 " shape=(None, 224, 224, 3),\n"
268 " dtype=tf.float32,\n"
269 " )\n"
270 " ],\n"
271 ")"
272 )
273 setattr(self, name, decorated_fn)
274 self._endpoint_names.append(name)
276 def add_variable_collection(self, name, variables):
277 """Register a set of variables to be retrieved after reloading.
279 Arguments:
280 name: The string name for the collection.
281 variables: A tuple/list/set of `tf.Variable` instances.
283 Example:
285 ```python
286 export_archive = ExportArchive()
287 export_archive.track(model)
288 # Register an endpoint
289 export_archive.add_endpoint(
290 name="serve",
291 fn=model.call,
292 input_signature=[tf.TensorSpec(shape=(None, 3), dtype=tf.float32)],
293 )
294 # Save a variable collection
295 export_archive.add_variable_collection(
296 name="optimizer_variables", variables=model.optimizer.variables)
297 export_archive.write_out("path/to/location")
299 # Reload the object
300 revived_object = tf.saved_model.load("path/to/location")
301 # Retrieve the variables
302 optimizer_variables = revived_object.optimizer_variables
303 ```
304 """
305 if not isinstance(variables, (list, tuple, set)):
306 raise ValueError(
307 "Expected `variables` to be a list/tuple/set. "
308 f"Received instead object of type '{type(variables)}'."
309 )
310 if not all(isinstance(v, tf.Variable) for v in variables):
311 raise ValueError(
312 "Expected all elements in `variables` to be "
313 "`tf.Variable` instances. Found instead the following types: "
314 f"{list(set(type(v) for v in variables))}"
315 )
316 setattr(self, name, list(variables))
318 def write_out(self, filepath, options=None):
319 """Write the corresponding SavedModel to disk.
321 Arguments:
322 filepath: `str` or `pathlib.Path` object.
323 Path where to save the artifact.
324 options: `tf.saved_model.SaveOptions` object that specifies
325 SavedModel saving options.
327 **Note on TF-Serving**: all endpoints registered via `add_endpoint()`
328 are made visible for TF-Serving in the SavedModel artifact. In addition,
329 the first endpoint registered is made visible under the alias
330 `"serving_default"` (unless an endpoint with the name
331 `"serving_default"` was already registered manually),
332 since TF-Serving requires this endpoint to be set.
333 """
334 if not self._endpoint_names:
335 raise ValueError(
336 "No endpoints have been set yet. Call add_endpoint()."
337 )
338 self._filter_and_track_resources()
340 signatures = {}
341 for name in self._endpoint_names:
342 signatures[name] = self._get_concrete_fn(name)
343 # Add "serving_default" signature key for TFServing
344 if "serving_default" not in self._endpoint_names:
345 signatures["serving_default"] = self._get_concrete_fn(
346 self._endpoint_names[0]
347 )
348 tf.saved_model.save(
349 self, filepath, options=options, signatures=signatures
350 )
351 # Print out available endpoints
352 endpoints = "\n\n".join(
353 _print_signature(getattr(self, name), name)
354 for name in self._endpoint_names
355 )
356 io_utils.print_msg(
357 f"Saved artifact at '{filepath}'. "
358 "The following endpoints are available:\n\n"
359 f"{endpoints}"
360 )
362 def _get_concrete_fn(self, endpoint):
363 """Workaround for some SavedModel quirks."""
364 if endpoint in self._endpoint_signatures:
365 return getattr(self, endpoint)
366 else:
367 traces = getattr(self, endpoint)._trackable_children("saved_model")
368 return list(traces.values())[0]
370 def _get_variables_used_by_endpoints(self):
371 fns = [self._get_concrete_fn(name) for name in self._endpoint_names]
372 return _list_variables_used_by_fns(fns)
374 def _filter_and_track_resources(self):
375 """Track resources used by endpoints / referenced in `track()` calls."""
376 # Start by extracting variables from endpoints.
377 fns = [self._get_concrete_fn(name) for name in self._endpoint_names]
378 tvs, ntvs = _list_variables_used_by_fns(fns)
379 self._all_variables = list(tvs + ntvs)
381 # Next, track lookup tables.
382 # Hopefully, one day this will be automated at the tf.function level.
383 self._misc_assets = []
384 from keras.src.layers.preprocessing.index_lookup import IndexLookup
386 if hasattr(self, "_tracked"):
387 for root in self._tracked:
388 descendants = tf.train.TrackableView(root).descendants()
389 for trackable in descendants:
390 if isinstance(trackable, IndexLookup):
391 self._misc_assets.append(trackable)
394def export_model(model, filepath):
395 export_archive = ExportArchive()
396 export_archive.track(model)
397 if isinstance(model, (functional.Functional, sequential.Sequential)):
398 input_signature = tf.nest.map_structure(_make_tensor_spec, model.inputs)
399 if isinstance(input_signature, list) and len(input_signature) > 1:
400 input_signature = [input_signature]
401 export_archive.add_endpoint("serve", model.__call__, input_signature)
402 else:
403 save_spec = model._get_save_spec()
404 if not save_spec:
405 raise ValueError(
406 "The model provided has never called. "
407 "It must be called at least once before export."
408 )
409 input_signature = [save_spec]
410 export_archive.add_endpoint("serve", model.__call__, input_signature)
411 export_archive.write_out(filepath)
414class ReloadedLayer(base_layer.Layer):
415 """Reload a Keras model/layer that was saved via SavedModel / ExportArchive.
417 Arguments:
418 filepath: `str` or `pathlib.Path` object. The path to the SavedModel.
419 call_endpoint: Name of the endpoint to use as the `call()` method
420 of the reloaded layer. If the SavedModel was created
421 via `model.export()`,
422 then the default endpoint name is `'serve'`. In other cases
423 it may be named `'serving_default'`.
425 Example:
427 ```python
428 model.export("path/to/artifact")
429 reloaded_layer = ReloadedLayer("path/to/artifact")
430 outputs = reloaded_layer(inputs)
431 ```
433 The reloaded object can be used like a regular Keras layer, and supports
434 training/fine-tuning of its trainable weights. Note that the reloaded
435 object retains none of the internal structure or custom methods of the
436 original object -- it's a brand new layer created around the saved
437 function.
439 **Limitations:**
441 * Only call endpoints with a single `inputs` tensor argument
442 (which may optionally be a dict/tuple/list of tensors) are supported.
443 For endpoints with multiple separate input tensor arguments, consider
444 subclassing `ReloadedLayer` and implementing a `call()` method with a
445 custom signature.
446 * If you need training-time behavior to differ from inference-time behavior
447 (i.e. if you need the reloaded object to support a `training=True` argument
448 in `__call__()`), make sure that the training-time call function is
449 saved as a standalone endpoint in the artifact, and provide its name
450 to the `ReloadedLayer` via the `call_training_endpoint` argument.
451 """
453 def __init__(
454 self,
455 filepath,
456 call_endpoint="serve",
457 call_training_endpoint=None,
458 trainable=True,
459 name=None,
460 dtype=None,
461 ):
462 # Initialize an empty layer, then add_weight() etc. as needed.
463 super().__init__(trainable=trainable, name=name, dtype=dtype)
465 self._reloaded_obj = tf.saved_model.load(filepath)
467 self.filepath = filepath
468 self.call_endpoint = call_endpoint
469 self.call_training_endpoint = call_training_endpoint
471 # Resolve the call function.
472 if hasattr(self._reloaded_obj, call_endpoint):
473 # Case 1: it's set as an attribute.
474 self.call_endpoint_fn = getattr(self._reloaded_obj, call_endpoint)
475 elif call_endpoint in self._reloaded_obj.signatures:
476 # Case 2: it's listed in the `signatures` field.
477 self.call_endpoint_fn = self._reloaded_obj.signatures[call_endpoint]
478 else:
479 raise ValueError(
480 f"The endpoint '{call_endpoint}' is neither an "
481 "attribute of the reloaded SavedModel, nor an entry "
482 "in the `signatures` field of the reloaded SavedModel. "
483 )
485 # Resolving the training function.
486 if call_training_endpoint:
487 if hasattr(self._reloaded_obj, call_training_endpoint):
488 self.call_training_endpoint_fn = getattr(
489 self._reloaded_obj, call_training_endpoint
490 )
491 elif call_training_endpoint in self._reloaded_obj.signatures:
492 self.call_training_endpoint_fn = self._reloaded_obj.signatures[
493 call_training_endpoint
494 ]
495 else:
496 raise ValueError(
497 f"The endpoint '{call_training_endpoint}' is "
498 "neither an attribute of the reloaded SavedModel, "
499 "nor an entry in the `signatures` field of "
500 "the reloaded SavedModel. "
501 )
503 # Add trainable and non-trainable weights from the call_endpoint_fn.
504 all_fns = [self.call_endpoint_fn]
505 if call_training_endpoint:
506 all_fns.append(self.call_training_endpoint_fn)
507 tvs, ntvs = _list_variables_used_by_fns(all_fns)
508 for v in tvs:
509 self._add_existing_weight(v, trainable=True)
510 for v in ntvs:
511 self._add_existing_weight(v, trainable=False)
512 self.built = True
514 def _add_existing_weight(self, weight, trainable):
515 """Calls add_weight() to register but not create an existing weight."""
516 self.add_weight(
517 name=weight.name,
518 shape=weight.shape,
519 dtype=weight.dtype,
520 trainable=trainable,
521 getter=lambda *_, **__: weight,
522 )
524 def call(self, inputs, training=False, **kwargs):
525 if training:
526 if self.call_training_endpoint:
527 return self.call_training_endpoint_fn(inputs, **kwargs)
528 return self.call_endpoint_fn(inputs, **kwargs)
530 def get_config(self):
531 base_config = super().get_config()
532 config = {
533 # Note: this is not intended to be portable.
534 "filepath": self.filepath,
535 "call_endpoint": self.call_endpoint,
536 "call_training_endpoint": self.call_training_endpoint,
537 }
538 return {**base_config, **config}
541def _make_tensor_spec(x):
542 return tf.TensorSpec(x.shape, dtype=x.dtype, name=x.name)
545def _print_signature(fn, name):
546 concrete_fn = fn._list_all_concrete_functions()[0]
547 pprinted_signature = concrete_fn.pretty_printed_signature(verbose=True)
548 lines = pprinted_signature.split("\n")
549 lines = [f"* Endpoint '{name}'"] + lines[1:]
550 endpoint = "\n".join(lines)
551 return endpoint
554def _list_variables_used_by_fns(fns):
555 trainable_variables = []
556 non_trainable_variables = []
557 trainable_variables_ids = set()
558 non_trainable_variables_ids = set()
559 for fn in fns:
560 if hasattr(fn, "concrete_functions"):
561 concrete_functions = fn.concrete_functions
562 elif hasattr(fn, "get_concrete_function"):
563 concrete_functions = [fn.get_concrete_function()]
564 else:
565 concrete_functions = [fn]
566 for concrete_fn in concrete_functions:
567 for v in concrete_fn.trainable_variables:
568 if id(v) not in trainable_variables_ids:
569 trainable_variables.append(v)
570 trainable_variables_ids.add(id(v))
572 for v in concrete_fn.variables:
573 if (
574 id(v) not in trainable_variables_ids
575 and id(v) not in non_trainable_variables_ids
576 ):
577 non_trainable_variables.append(v)
578 non_trainable_variables_ids.add(id(v))
579 return trainable_variables, non_trainable_variables