# Copyright 2018 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. # ============================================================================== """Tools for deserializing `Function`s.""" import collections import pprint import re from absl import logging from tensorflow.core.function import trace_type from tensorflow.core.function.polymorphism import function_type as function_type_lib from tensorflow.core.protobuf import saved_object_graph_pb2 from tensorflow.python.eager import def_function from tensorflow.python.eager import function as function_lib from tensorflow.python.eager.polymorphic_function import function_type_utils from tensorflow.python.framework import func_graph as func_graph_lib from tensorflow.python.framework import function_def_to_graph as function_def_lib from tensorflow.python.framework import op_def_registry from tensorflow.python.framework import ops from tensorflow.python.framework import tensor from tensorflow.python.framework import type_spec from tensorflow.python.ops import array_ops from tensorflow.python.ops import custom_gradient from tensorflow.python.ops import default_gradient from tensorflow.python.ops import resource_variable_ops from tensorflow.python.saved_model import nested_structure_coder from tensorflow.python.util import compat from tensorflow.python.util import nest from tensorflow.python.util import tf_decorator from tensorflow.python.util import tf_inspect def _is_tensor(t): return isinstance( t, (tensor.Tensor, resource_variable_ops.BaseResourceVariable)) # TODO(b/205016027): Update this to just use ConcreteFunction.__call__ with the # structured signature. def _call_concrete_function(function, inputs): """Calls a restored Function with structured inputs. This differs from `function.__call__` in that inputs and outputs are structured and that it casts inputs to tensors if needed. Note: this does not checks that non-tensor inputs match. That should be done before via `_concrete_function_callable_with`. Args: function: ConcreteFunction to call. inputs: Structured inputs compatible with `function.graph.structured_input_signature`. Returns: The structured function output. """ expected_structure = function.graph.structured_input_signature flatten_inputs = nest.flatten_up_to( expected_structure, inputs, expand_composites=True) flatten_expected = nest.flatten(expected_structure, expand_composites=True) tensor_inputs = [] for arg, expected in zip(flatten_inputs, flatten_expected): if isinstance(expected, tensor.TensorSpec): tensor_inputs.append( ops.convert_to_tensor(arg, dtype_hint=expected.dtype)) elif isinstance(expected, resource_variable_ops.VariableSpec): tensor_inputs.append(arg.handle) result = function._call_flat(tensor_inputs, function.captured_inputs) # pylint: disable=protected-access if isinstance(result, ops.Operation): return None return result def _try_convert_to_tensor_spec(arg, dtype_hint): """Returns None or TensorSpec obtained if `arg` is converted to tensor.""" try: # Note: try conversion in a FuncGraph to avoid polluting current context. with func_graph_lib.FuncGraph(name="guess_conversion").as_default(): result = ops.convert_to_tensor(arg, dtype_hint=dtype_hint) return tensor.TensorSpec(shape=result.shape, dtype=result.dtype) except (TypeError, ValueError): return None def _concrete_function_callable_with(function, inputs, allow_conversion): """Returns whether concrete `function` can be called with `inputs`.""" expected_structure = function.graph.structured_input_signature try: flatten_inputs = nest.flatten_up_to(expected_structure, inputs) except (TypeError, ValueError): return False for arg, expected in zip(flatten_inputs, nest.flatten(expected_structure)): if isinstance(expected, tensor.TensorSpec): if allow_conversion: arg = _try_convert_to_tensor_spec(arg, dtype_hint=expected.dtype) if not _is_tensor(arg) and not isinstance(arg, tensor.TensorSpec): return False if arg.dtype != expected.dtype: return False if not expected.shape.is_compatible_with(arg.shape): return False elif isinstance(expected, type_spec.TypeSpec): if not expected.is_compatible_with(arg): return False elif _is_tensor(arg): if id(arg) != id(expected): return False else: if arg != expected: return False return True def _deserialize_function_spec_as_nonmethod(function_spec_proto): """Deserialize a FunctionSpec object from its proto representation.""" typeless_fullargspec = nested_structure_coder.decode_proto( function_spec_proto.fullargspec) # Convert a method function into a non method. if function_spec_proto.is_method or ( typeless_fullargspec.args and typeless_fullargspec.args[0] == "self" ): if not typeless_fullargspec.args: raise NotImplementedError( "Cannot deserialize a method function without a named " "'self' argument.") args = typeless_fullargspec.args[1:] else: args = typeless_fullargspec.args fullargspec = tf_inspect.FullArgSpec( args=args, varargs=typeless_fullargspec.varargs, varkw=typeless_fullargspec.varkw, defaults=typeless_fullargspec.defaults, kwonlyargs=typeless_fullargspec.kwonlyargs, kwonlydefaults=typeless_fullargspec.kwonlydefaults, annotations=typeless_fullargspec.annotations) input_signature = nested_structure_coder.decode_proto( function_spec_proto.input_signature) # See `tf.function` and the JitCompile proto for details. jit_compile = { saved_object_graph_pb2.FunctionSpec.JitCompile.DEFAULT: None, saved_object_graph_pb2.FunctionSpec.JitCompile.ON: True, saved_object_graph_pb2.FunctionSpec.JitCompile.OFF: False, }.get(function_spec_proto.jit_compile) return function_type_utils.FunctionSpec.from_fullargspec_and_signature( fullargspec=fullargspec, input_signature=input_signature, jit_compile=jit_compile) # TODO(b/203440205): Set FunctionType with ConcreteFunction constructor. def set_preinitialized_function_spec(concrete_fn, spec): """Set the FunctionType of the ConcreteFunction using FunctionSpec.""" if spec is None: concrete_fn._function_type = None # pylint: disable=protected-access return unconstrained_type = function_type_lib.FunctionType( [ function_type_lib.Parameter(p.name, p.kind, p.optional, None) for p in spec.function_type.parameters.values() ] ) arg_specs, kwarg_specs = concrete_fn.structured_input_signature input_function_type, _ = function_type_lib.canonicalize_to_monomorphic( arg_specs, { function_type_lib.sanitize_arg_name(k): v for k, v in kwarg_specs.items() }, spec.default_values, {}, unconstrained_type, ) output_type = trace_type.from_value(concrete_fn.graph.structured_outputs) # Captures are restored later so we will update it then. function_type = function_type_lib.FunctionType( input_function_type.parameters.values(), return_annotation=output_type, ) concrete_fn._function_type = function_type # pylint: disable=protected-access # TODO(b/205016761): The fact that we can't derive ConcreteFunction calling # conventions from the serialized input spec right now is unfortunate. Merging # these would be good, maybe by adding TensorSpec names to cache keys so renamed # keyword arguments would yield different ConcreteFunctions. def setup_bare_concrete_function(saved_bare_concrete_function, concrete_functions): """Makes a restored bare concrete function callable.""" concrete_function = concrete_functions[ saved_bare_concrete_function.concrete_function_name] # pylint: disable=protected-access concrete_function._arg_keywords = ( saved_bare_concrete_function.argument_keywords) concrete_function._num_positional_args = ( saved_bare_concrete_function.allowed_positional_arguments) if saved_bare_concrete_function.HasField("function_spec"): function_spec = _deserialize_function_spec_as_nonmethod( saved_bare_concrete_function.function_spec) set_preinitialized_function_spec(concrete_function, function_spec) # pylint: enable=protected-access concrete_function.add_to_graph() return concrete_function class RestoredFunction(def_function.Function): """Wrapper class for a function that has been restored from saved state. See `def_function.Function`. """ def __init__(self, python_function, name, function_spec, concrete_functions): # TODO(b/205016819): We may enable autograph once exceptions are supported. super(RestoredFunction, self).__init__( python_function, name, autograph=False, jit_compile=function_spec.jit_compile) self.concrete_functions = concrete_functions self._function_type = function_spec.function_type self._default_values = function_spec.default_values # Prevent RestoredFunction from spamming users with frequent tracing # warnings. self._omit_frequent_tracing_warning = True @property def _run_functions_eagerly(self): # We do not have access to the original python function, and thus, we # cannot meaningfully do anything but call our concrete function graphs # under the hood. # # Attempting to call our bespoke python function (i.e. # `restored_function_body`) will work so long as the user passes in all # required and optional arguments. If an optional argument is missing, # however, the call will break. For this reason, we instead skip the # eager call path altogether if a user has enabled eager function execution # via `tf.config.run_functions_eagerly`. return False def _list_all_concrete_functions(self): return self.concrete_functions def _list_all_concrete_functions_for_serialization(self): return self.concrete_functions def recreate_function(saved_function, concrete_functions): """Creates a `Function` from a `SavedFunction`. Args: saved_function: `SavedFunction` proto. concrete_functions: map from function name to `ConcreteFunction`. As a side effect of this function, the `FunctionSpec` from `saved_function` is added to each `ConcreteFunction` in this map. Returns: A `Function`. """ # TODO(b/205017389): Construct a `Function` with the cache populated # instead of creating a new `Function` backed by a Python layer to # glue things together. Current approach is nesting functions deeper for each # serialization cycle. # Note: handling method functions is tricky since make_decorator does not # allows control of "ismethod". Additionally since restored functions do # not behave as methods i.e. they always use the same captured tensors # independent of the object they are bound to, there is little value on # propagating that correctly. # # Ideally this conversion should happen at serialization time. But since # there are SavedModels which have "ismethod" populated and have an extra # argument that they expect to be ignored, we do it at deserialization. function_spec = _deserialize_function_spec_as_nonmethod( saved_function.function_spec) def restored_function_body(*args, **kwargs): """Calls a restored function or raises an error if no matching function.""" if not saved_function.concrete_functions: raise ValueError("Found zero restored functions for caller function.") # This is the format of function.graph.structured_input_signature. At this # point, the args and kwargs have already been canonicalized. inputs = (args, kwargs) # First try to find a concrete function that can be called without input # conversions. This allows one to pick a more specific trace in case there # was also a more expensive one that supported tensors. for allow_conversion in [False, True]: for function_name in saved_function.concrete_functions: function = concrete_functions[function_name] if any([inp is None for inp in function.captured_inputs]): raise ValueError("Looks like you are trying to run a loaded " "non-Keras model that was trained using " "tf.distribute.experimental.ParameterServerStrategy " "with variable partitioning, which is not currently " "supported. Try using Keras to define your model " "if possible.") if _concrete_function_callable_with(function, inputs, allow_conversion): return _call_concrete_function(function, inputs) signature_descriptions = [] def _pretty_format_positional(positional): return "Positional arguments ({} total):\n * {}".format( len(positional), "\n * ".join(pprint.pformat(a) for a in positional)) for index, function_name in enumerate(saved_function.concrete_functions): concrete_function = concrete_functions[function_name] positional, keyword = concrete_function.structured_input_signature signature_descriptions.append( "Option {}:\n {}\n Keyword arguments: {}".format( index + 1, _pretty_format_positional(positional), keyword)) raise ValueError( "Could not find matching concrete function to call loaded from the " f"SavedModel. Got:\n {_pretty_format_positional(args)}\n Keyword " f"arguments: {kwargs}\n\n Expected these arguments to match one of the " f"following {len(saved_function.concrete_functions)} option(s):\n\n" f"{(chr(10)+chr(10)).join(signature_descriptions)}") concrete_function_objects = [] for concrete_function_name in saved_function.concrete_functions: concrete_function_objects.append(concrete_functions[concrete_function_name]) for cf in concrete_function_objects: set_preinitialized_function_spec(cf, function_spec) restored_function = RestoredFunction(restored_function_body, restored_function_body.__name__, function_spec, concrete_function_objects) return tf_decorator.make_decorator( restored_function_body, restored_function, decorator_argspec=function_spec.fullargspec) def load_function_def_library(library, saved_object_graph=None, load_shared_name_suffix=None, wrapper_function=None): """Load a set of functions as concrete functions without captured inputs. Functions names are manipulated during load such that they do not overlap with previously created ones. Gradients are re-registered under new names. Ops that reference the gradients are updated to reflect the new registered names. Args: library: FunctionDefLibrary proto message. saved_object_graph: SavedObjectGraph proto message. If not passed in, concrete function structured signatures and outputs will not be set. load_shared_name_suffix: If specified, used to uniquify shared names. Otherwise, a unique name is generated. wrapper_function: An object that will be wrapped on newly created functions. Returns: Map of original function names in the library to instances of `ConcreteFunction` without captured inputs. Raises: ValueError: if functions dependencies have a cycle. """ library_function_names = set(fdef.signature.name for fdef in library.function) functions = {} renamed_functions = {} # Our graph building code currently requires functions to be registered with # some tf.Graph in order to import functions using the # op-name-is-function-name calling convention. To avoid leaking memory into # the global default graph when executing eagerly, we create a temporary # Graph. # # TODO(b/205023033): Make this Graph creation unnecessary when executing # eagerly by fixing function_def_to_graph_def. if ops.executing_eagerly_outside_functions(): graph = ops.Graph() else: graph = ops.get_default_graph() if load_shared_name_suffix is None: load_shared_name_suffix = "_load_{}".format(ops.uid()) # Custom gradient functions must be re-registered under new UIDs. library_gradient_names = {} # Maps old op type to old function name new_gradient_op_types = {} # Maps old gradient op type to new op type. gradients_to_register = {} # Maps old function name to new op type for gdef in library.registered_gradients: if gdef.registered_op_type: new_op_type = custom_gradient.generate_name() old_op_type = compat.as_bytes(gdef.registered_op_type) library_gradient_names[old_op_type] = gdef.gradient_func new_gradient_op_types[old_op_type] = new_op_type gradients_to_register[gdef.gradient_func] = new_op_type function_deps = {} for fdef in library.function: function_deps[fdef.signature.name] = _list_function_deps( fdef, library_function_names, library_gradient_names) loaded_gradients = {} for fdef in _sort_function_defs(library, function_deps): orig_name = _fix_fdef_in_place(fdef, functions, load_shared_name_suffix, new_gradient_op_types) # Setup function signatures and outputs # # When concrete functions are created normally (i.e. when they're originally # created and not loaded via saved model), the inputs and outputs are # calculated based on the values passed in by the user and returned from the # original function, respectively. We don't have access to those anymore at # restore time, so we must instead pass them to the FuncGraph explicitly. structured_input_signature = None structured_outputs = None if (saved_object_graph is not None and orig_name in saved_object_graph.concrete_functions): # TODO(b/204324043): Offload the deserialization of the protos to the # first class objects by passing the actual protos. This is blocked on # importing `nested_structure_coder` in function.py causing a circular # dependency. proto = saved_object_graph.concrete_functions[orig_name] structured_input_signature = nested_structure_coder.decode_proto( proto.canonicalized_input_signature) structured_outputs = nested_structure_coder.decode_proto( proto.output_signature) # There is no need to copy all functions into the function def graph. It # leads to a O(n^2) increase of memory when importing functions and the # extra function definitions are a no-op since they already imported as a # function before and passed in explicitly (due to the topologic sort # import). with graph.as_default(): func_graph = function_def_lib.function_def_to_graph( fdef, structured_input_signature=structured_input_signature, structured_outputs=structured_outputs) # Restores gradients for function-call ops (not the same as ops that use # custom gradients) _restore_gradient_functions(func_graph, renamed_functions, loaded_gradients) for dep in function_deps[orig_name]: functions[dep].add_to_graph(func_graph) # We do not initialize the new ConcreteFunction's function_spec and/or # arg_keywords here (which are used to parse the structured and flat # signatures, respectively). ConcreteFunction that are part of a saved # function is set up later by recreate_function(); and bare ConcreteFunction # is set up by by setup_bare_concrete_function(). # However, we copy the FunctionDef attributes to the new ConcreteFunction, # excluding the "_input_shapes", which may cause an error during input shape # initialization at a later stage. if "_input_shapes" in fdef.attr: del fdef.attr["_input_shapes"] function_type = function_type_lib.from_structured_signature( func_graph.structured_input_signature, func_graph.structured_outputs, func_graph.function_captures.capture_types, ) func = function_lib.ConcreteFunction.from_func_graph( func_graph, function_type, attrs=fdef.attr) if wrapper_function: func = wrapper_function(func) func.add_to_graph(graph) functions[orig_name] = func renamed_functions[func.name] = func if any(op.type == "TRTEngineOp" for op in func_graph.get_operations()): # TODO(b/150708051): Remove this hack once TensorRT SavedModel integration # is fixed. Currently it's leaking memory to maintain bug compatibility # with previous behavior. func.add_to_graph(ops.get_default_graph()) if orig_name in gradients_to_register: gradient_op_type = gradients_to_register[orig_name] loaded_gradients[compat.as_bytes(gradient_op_type)] = func ops.RegisterGradient(gradient_op_type)(_gen_gradient_func(func)) return functions def _gen_gradient_func(func): """Wraps a deserialized function.""" def gradient_func(unused_op, *result_grads): # Replace all `None` arguments, because the traced custom gradient function # expects tensors. Replacing with zeros is correct since the `None` values # occur when the gradient is unconnected, and thus the gradient is # "statically proven to be zero." See `tf.UnconnectedGradients` for details. def none_to_zero(x, t): if x is not None: return x shape, dtype = default_gradient.shape_and_dtype(t) if shape.is_fully_defined(): return default_gradient.zeros_like(t) dims = [] if shape.rank is not None: dims = [1 if d is None else d for d in shape.as_list()] return array_ops.zeros(dims, dtype) result_grads = [ none_to_zero(x, t) for (x, t) in zip(result_grads, func.graph.inputs) ] return func(*result_grads) return gradient_func def _restore_gradient_functions(func_graph, renamed_functions, loaded_gradients): """Populate function op's _gradient_function with default gradient.""" for op in func_graph.get_operations(): # TODO(b/205024208): This code assumes that the gradient registered for this # function call is the default gradient for the function and not a custom # one. if op.type in ["StatefulPartitionedCall", "PartitionedCall"]: function = renamed_functions[compat.as_bytes( op.node_def.attr["f"].func.name)] op._gradient_function = function._get_gradient_function() # pylint: disable=protected-access try: gradient_op_type = op.get_attr("_gradient_op_type") except ValueError: pass else: if gradient_op_type in loaded_gradients: grad_fn = loaded_gradients[gradient_op_type] grad_fn._num_positional_args = len(op.inputs) # pylint: disable=protected-access grad_fn._arg_keywords = [inp.name for inp in op.inputs] # pylint: disable=protected-access def _sort_function_defs(library, function_deps): """Return a topologic sort of FunctionDefs in a library.""" edges = collections.defaultdict(list) in_count = collections.defaultdict(lambda: 0) for fname, deps in function_deps.items(): for dep in deps: edges[dep].append(fname) in_count[fname] += 1 ready = [ fdef.signature.name for fdef in library.function if in_count[fdef.signature.name] == 0 ] output = [] while ready: node = ready.pop() output.append(node) for dest in edges[node]: in_count[dest] -= 1 if not in_count[dest]: ready.append(dest) if len(output) != len(library.function): failed_to_resolve = sorted(set(in_count.keys()) - set(output)) raise ValueError("There is a cyclic dependency between functions. ", f"Could not resolve {failed_to_resolve}.") reverse = {fdef.signature.name: fdef for fdef in library.function} return [reverse[x] for x in output] def _get_gradient_op_type(node_def): """Returns the custom gradient op type.""" if ("_gradient_op_type" in node_def.attr and node_def.op not in ["StatefulPartitionedCall", "PartitionedCall"]): return node_def.attr["_gradient_op_type"].s return None def fix_node_def(node_def, functions, shared_name_suffix): """Replace functions calls and shared names in `node_def`.""" if node_def.op in functions: node_def.op = functions[node_def.op].name for _, attr_value in node_def.attr.items(): if attr_value.WhichOneof("value") == "func": attr_value.func.name = functions[attr_value.func.name].name elif attr_value.WhichOneof("value") == "list": for fn in attr_value.list.func: fn.name = functions[fn.name].name # Fix old table creation bug. if node_def.op == "HashTableV2": if ("use_node_name_sharing" not in node_def.attr or not node_def.attr["use_node_name_sharing"].b): node_def.attr["use_node_name_sharing"].b = True # We are turning on node mame sharing, so have to make sure we don't # accidentally share a table resource. shared_name_suffix += "_{}".format(ops.uid()) # TODO(b/124205571): Avoid accidental sharing and destruction of restored # resources. For now uniquify "shared_name" when loading functions to avoid # sharing. # TODO: Add regression test for b/150826922. op_def = op_def_registry.get(node_def.op) if op_def: attr = next((a for a in op_def.attr if a.name == "shared_name"), None) if attr: shared_name = None if "shared_name" in node_def.attr and node_def.attr["shared_name"].s: shared_name = node_def.attr["shared_name"].s elif attr.default_value.s: shared_name = compat.as_bytes(attr.default_value.s) if not shared_name: shared_name = compat.as_bytes(node_def.name) node_def.attr["shared_name"].s = ( shared_name + compat.as_bytes(shared_name_suffix)) def _fix_fdef_in_place(fdef, functions, shared_name_suffix, new_gradient_op_types): """Fixes a FunctionDef proto to be loaded in current context. In particular, when loading a function library into an eager context, one must rename the functions to avoid conflicts with existent functions. Args: fdef: FunctionDef proto to fix. It is mutated in-place. functions: map from function name to a ConcreteFunction instance. shared_name_suffix: A unique string for this load which helps to avoid `shared_name` collisions across loads. Two functions from the same load using the same `shared_name` still need to share, but functions from different loads with the same `shared_name` should not. new_gradient_op_types: map from old gradient op type to newly generated op type. Returns: orig_name: original value of fdef.signature.name """ orig_name = fdef.signature.name contains_unsaved_custom_gradients = False for node_def in fdef.node_def: fix_node_def(node_def, functions, shared_name_suffix) op_type = _get_gradient_op_type(node_def) if op_type is not None: if op_type in new_gradient_op_types: node_def.attr["_gradient_op_type"].s = compat.as_bytes( new_gradient_op_types[op_type]) else: contains_unsaved_custom_gradients = True if contains_unsaved_custom_gradients: logging.warning( "Importing a function (%s) with ops with unsaved custom gradients. Will" " likely fail if a gradient is requested.", fdef.signature.name) fdef.signature.name = _clean_function_name(fdef.signature.name) return orig_name def _list_function_deps(fdef, library_function_names, library_gradient_names): """Find functions referenced in `fdef`.""" # TODO(b/205023953): Recurse into list attributes and into NameAttrList attrs # both when listing deps and when fixing them. `function_def_to_graph` also # requires fixes. deps = set() for node_def in fdef.node_def: grad_op_type = _get_gradient_op_type(node_def) if node_def.op in library_function_names: deps.add(node_def.op) elif grad_op_type and grad_op_type in library_gradient_names: deps.add(library_gradient_names[grad_op_type]) else: for _, attr_value in node_def.attr.items(): if attr_value.WhichOneof("value") == "func": deps.add(attr_value.func.name) elif attr_value.WhichOneof("value") == "list": for fn in attr_value.list.func: deps.add(fn.name) return deps _FUNCTION_WRAPPER_NAME_REGEX = r"^%s(.*)_\d+$" % (function_lib._INFERENCE_PREFIX ) # pylint:disable=protected-access def _clean_function_name(name): """Vanity function to keep the function names comprehensible.""" # Note: each time a function is wrapped into `function_lib.ConcreteFunction` # its name becomes "__inference__xyz". match = re.search(_FUNCTION_WRAPPER_NAME_REGEX, name) if match: return match.group(1) else: return name