wehub-resource-sync
73 lines
2.6 KiB
Python
73 lines
2.6 KiB
Python
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Tracing utilities used by SavedModel."""
|
|
|
|
from tensorflow.python.checkpoint import saveable_compat
|
|
from tensorflow.python.checkpoint import tensor_callable
|
|
from tensorflow.python.eager import def_function
|
|
from tensorflow.python.eager import function as defun
|
|
|
|
|
|
def trace_save_and_restore(obj):
|
|
"""Traces `Trackable` serialize- and restore-from-tensors functions.
|
|
|
|
Args:
|
|
obj: A `Trackable` object.
|
|
|
|
Returns:
|
|
A concrete Function.
|
|
"""
|
|
legacy_name = saveable_compat.get_saveable_name(obj)
|
|
|
|
obj_save_fn = obj._serialize_to_tensors # pylint: disable=protected-access
|
|
obj_restore_fn = obj._restore_from_tensors # pylint: disable=protected-access
|
|
|
|
if isinstance(obj_save_fn, defun.ConcreteFunction):
|
|
concrete_save = obj_save_fn
|
|
else:
|
|
@def_function.function
|
|
def save_fn():
|
|
tensor_dict = obj_save_fn()
|
|
if any(isinstance(v, tensor_callable.Callable)
|
|
for v in tensor_dict.values()):
|
|
raise NotImplementedError(
|
|
f"Unable to export SavedModel with object of type {type(obj)} "
|
|
"because it returns a Callable in `_serialize_to_tensors`. "
|
|
"If you need this functionality please file a feature request.")
|
|
|
|
if legacy_name:
|
|
# If there is a legacy decorator, append the name to the keys.
|
|
return {f"{legacy_name}{key}": value
|
|
for key, value in tensor_dict.items()}
|
|
return tensor_dict
|
|
|
|
concrete_save = save_fn.get_concrete_function()
|
|
|
|
if isinstance(obj_restore_fn, defun.ConcreteFunction):
|
|
concrete_restore = obj_restore_fn
|
|
else:
|
|
@def_function.function
|
|
def restore_fn(restored_tensors):
|
|
if legacy_name:
|
|
# Do the opposite operation of save_fn()
|
|
restored_tensors = {key[len(legacy_name):]: value
|
|
for key, value in restored_tensors.items()}
|
|
obj_restore_fn(restored_tensors)
|
|
|
|
concrete_restore = restore_fn.get_concrete_function(
|
|
concrete_save.structured_outputs)
|
|
|
|
return concrete_save, concrete_restore
|