394 lines
16 KiB
Python
394 lines
16 KiB
Python
# 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.
|
|
# ==============================================================================
|
|
|
|
"""Utilities for V2 control flow."""
|
|
|
|
from tensorflow.core.framework import attr_value_pb2
|
|
from tensorflow.python.distribute import distribution_strategy_context
|
|
from tensorflow.python.eager import context
|
|
from tensorflow.python.eager import function
|
|
from tensorflow.python.framework import function_def_to_graph
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.framework.func_graph import FuncGraph
|
|
from tensorflow.python.ops import control_flow_util
|
|
from tensorflow.python.ops import control_flow_v2_func_graphs
|
|
from tensorflow.python.ops import gradients_util
|
|
from tensorflow.python.util import keras_deps
|
|
from tensorflow.python.util import tf_contextlib
|
|
|
|
|
|
_EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE = None
|
|
_DISABLE_LOWER_USING_SWITCH_MERGE = False
|
|
|
|
|
|
CondBranchFuncGraph = control_flow_v2_func_graphs.CondBranchFuncGraph
|
|
WhileCondFuncGraph = control_flow_v2_func_graphs.WhileCondFuncGraph
|
|
WhileBodyFuncGraph = control_flow_v2_func_graphs.WhileBodyFuncGraph
|
|
|
|
|
|
def in_defun():
|
|
"""Returns if the current graph is, or is nested in, a defun."""
|
|
if context.executing_eagerly(): return False
|
|
|
|
graph = ops.get_default_graph()
|
|
while (isinstance(graph, CondBranchFuncGraph) or
|
|
isinstance(graph, WhileBodyFuncGraph) or
|
|
isinstance(graph, WhileCondFuncGraph)):
|
|
graph = graph.outer_graph
|
|
return isinstance(graph, FuncGraph)
|
|
|
|
|
|
def in_while_loop_defun(graph):
|
|
"""Returns if the graph is a while loop FuncGraph."""
|
|
if context.executing_eagerly(): return False
|
|
return (isinstance(graph, WhileCondFuncGraph) or
|
|
isinstance(graph, WhileBodyFuncGraph))
|
|
|
|
|
|
def create_new_tf_function(func_graph):
|
|
"""Converts func_graph to a TF_Function and adds it to the current graph.
|
|
|
|
Args:
|
|
func_graph: FuncGraph
|
|
|
|
Returns:
|
|
The name of the new TF_Function.
|
|
"""
|
|
func = function._EagerDefinedFunction( # pylint: disable=protected-access
|
|
func_graph.name, func_graph, func_graph.inputs, func_graph.outputs, {})
|
|
func.add_to_graph(func_graph.outer_graph)
|
|
return func_graph.name
|
|
|
|
|
|
def unique_fn_name(scope, name):
|
|
"""Returns a unique name to use for a control flow function.
|
|
|
|
Args:
|
|
scope: A name scope string.
|
|
name: An identifier for this function (e.g. "true", "body").
|
|
|
|
Returns:
|
|
A string, the name to use for the function.
|
|
"""
|
|
return ("%s%s_%s" % (scope, name, ops.uid())).replace("/", "_")
|
|
|
|
|
|
def unique_grad_fn_name(forward_name):
|
|
return "%s_grad_%s" % (forward_name, ops.uid())
|
|
|
|
|
|
def maybe_set_lowering_attr(op, lower_using_switch_merge=None):
|
|
"""Sets the flag to enable lowering on `op` if necessary.
|
|
|
|
Lowering allows cond_v2 and while_v2 to avoid some of the limitations of
|
|
Functions, allowing users to specify devices & colocation inside of cond_v2
|
|
and while_v2 input functions, and enabling non-strict evaluation & partial
|
|
pruning. This brings v2 control flow closer to feature parity with v1 control
|
|
flow.
|
|
|
|
However, we do not lower in the following cases:
|
|
- When the `If` or `While` ops are in the XLA context. Because it is easier
|
|
for XLA to apply its own optimizations when dealing with un-lowered
|
|
control flow operators than with low-level control flow primitives.
|
|
- When the eager execution context specifies the executor of functions to
|
|
be the single threaded executor (see context.function_executor_type()).
|
|
Because the single threaded executor does not support v1 control flow ops.
|
|
- When 'lower_using_switch_merge' is explicitly set to False.
|
|
|
|
Args:
|
|
op: An `If` or `While` Operation.
|
|
lower_using_switch_merge: Explicit value to lower or not (optional).
|
|
"""
|
|
if lower_using_switch_merge is not None:
|
|
# pylint: disable=protected-access
|
|
op._set_attr("_lower_using_switch_merge",
|
|
attr_value_pb2.AttrValue(b=lower_using_switch_merge))
|
|
# pylint: enable=protected-access
|
|
elif (not _DISABLE_LOWER_USING_SWITCH_MERGE and
|
|
not control_flow_util.GraphOrParentsInXlaContext(op.graph) and
|
|
context.context().function_call_options.executor_type !=
|
|
"SINGLE_THREADED_EXECUTOR"):
|
|
# pylint: disable=protected-access
|
|
op._set_attr("_lower_using_switch_merge", attr_value_pb2.AttrValue(b=True))
|
|
# pylint: enable=protected-access
|
|
|
|
|
|
def maybe_propagate_compile_time_consts_in_xla(op):
|
|
"""Tells XLA whether to propagate compile-time consts in the loop body.
|
|
|
|
This is needed to make compile time constants available to ops, for example
|
|
`max_num_elements` in `EmptyTensorList`, inside the loop body. Ideally this
|
|
would always be turned on, but that doesn't work with legacy functionalized
|
|
while_loops.
|
|
|
|
Args:
|
|
op: A `While` Operation.
|
|
"""
|
|
if control_flow_util.GraphOrParentsInXlaContext(op.graph):
|
|
# pylint: disable=protected-access
|
|
op._set_attr("_xla_propagate_compile_time_consts",
|
|
attr_value_pb2.AttrValue(b=True))
|
|
# pylint: enable=protected-access
|
|
|
|
|
|
def resource_input_index(tensor_name, input_names, node_defs, functions):
|
|
"""Returns the index of the input corresponding to `tensor_name`.
|
|
|
|
This method is used to find the corresponding index of an arbitrary resource
|
|
tensor in a function (the function could be a loop body). We assume that
|
|
resource handles are never created in functions, so that every resource
|
|
tensor can be traced back to a function input.
|
|
|
|
The awkward signature of this method is to make it work with both FuncGraphs
|
|
and FunctionDefs. This is so we can recurse on function call ops without
|
|
building the corresponding FuncGraph (note that even if a FuncGraph for a
|
|
FunctionDef already exists, the input/output/node names may have been
|
|
changed when the FuncGraph was serialized to the FunctionDef, which makes it
|
|
unusable with this algorithm).
|
|
|
|
Args:
|
|
tensor_name: the name of the resource tensor to be resolved to an input.
|
|
input_names: a list of the names of all inputs to the function.
|
|
node_defs: a dict mapping op name -> NodeDef for every op in the function.
|
|
functions: a dict mapping function name -> _EagerDefinedFunction.
|
|
|
|
Returns:
|
|
The index into input_names corresponding to `tensor_name`.
|
|
"""
|
|
while tensor_name not in input_names:
|
|
# FunctionDefs and graphs use different tensor naming conventions.
|
|
parts = tensor_name.split(":")
|
|
if len(parts) == 3:
|
|
op_name, _, output_idx = parts
|
|
elif len(parts) == 2:
|
|
op_name, output_idx = parts
|
|
else:
|
|
assert len(parts) == 1
|
|
op_name = parts[0]
|
|
output_idx = 0
|
|
tensor_name = "%s:%d" % (tensor_name, output_idx)
|
|
# Check again for cases where the tensor suffix (":0") is stripped out.
|
|
if tensor_name in input_names:
|
|
break
|
|
output_idx = int(output_idx)
|
|
node_def = node_defs[op_name]
|
|
|
|
def _extract_input_index(function_attribute_name):
|
|
func_name = node_def.attr[function_attribute_name].func.name
|
|
fdef = functions[func_name].definition
|
|
output_arg_name = fdef.signature.output_arg[output_idx].name
|
|
output_tensor_name = fdef.ret[output_arg_name]
|
|
return resource_input_index(
|
|
output_tensor_name, [arg.name for arg in fdef.signature.input_arg],
|
|
{ndef.name: ndef for ndef in fdef.node_def}, functions)
|
|
|
|
if node_def.op in ("Identity", "While"):
|
|
# Captured resources occur at the same index in the lists of inputs and
|
|
# outputs of a while or identity op. So we lookup the input of `tensor.op`
|
|
# at the same index as the index of `tensor` in the `tensor.op.outputs`.
|
|
tensor_name = node_def.input[output_idx]
|
|
elif node_def.op in ("PartitionedCall", "StatefulPartitionedCall"):
|
|
# Functions output any captured resource tensors used by their
|
|
# gradients. `tensor_name` is one of these outputs from a nested
|
|
# function call, so recursively find the corresponding input in the
|
|
# nested FunctionDef.
|
|
tensor_name = node_def.input[_extract_input_index("f")]
|
|
elif node_def.op in ("If", "StatelessIf"):
|
|
input_index = _extract_input_index("then_branch")
|
|
if input_index != _extract_input_index("else_branch"):
|
|
raise AssertionError(
|
|
("Expected cond branches ({} op) to each have the same "
|
|
"input->output mapping of resources.").format(node_def.op))
|
|
tensor_name = node_def.input[
|
|
# Ignore the `cond` input; the function inputs come after.
|
|
input_index + 1]
|
|
else:
|
|
# We assume there are no other ops types that will "forward" resource
|
|
# handles like this, so all other handles must have been created by the
|
|
# op. (Note that cond_v2 wraps resource handle outputs in optionals,
|
|
# which we'll end up accumulating).
|
|
raise ValueError("Taking gradient of a while loop which creates "
|
|
"a resource in its body is not supported: %s (%s)"
|
|
% (op_name, node_def.op))
|
|
|
|
return input_names.index(tensor_name)
|
|
|
|
|
|
@tf_contextlib.contextmanager
|
|
def clear_control_inputs():
|
|
"""Clears the control inputs but preserves the ControlFlowContext.
|
|
|
|
This is needed to preserve the XLAControlFlowControl when clearing
|
|
control inputs for the gradient accumulators in while_v2.
|
|
`ops.control_dependencies` does not allow that.
|
|
|
|
Yields:
|
|
A context manager in which the ops created will not have any control inputs
|
|
by default but the control flow context is the same.
|
|
"""
|
|
# pylint: disable=protected-access
|
|
control_flow_context = ops.get_default_graph()._get_control_flow_context()
|
|
with ops.control_dependencies(None):
|
|
ops.get_default_graph()._set_control_flow_context(control_flow_context)
|
|
yield
|
|
# pylint: enable=protected-access
|
|
|
|
|
|
def _is_tpu_strategy(strategy):
|
|
return (strategy is not None and
|
|
strategy.__class__.__name__.startswith("TPUStrategy"))
|
|
|
|
|
|
def _is_building_keras_layer():
|
|
# TODO(srbs): Remove this function when we no long support session with Keras.
|
|
keras_call_context_function = keras_deps.get_call_context_function()
|
|
if keras_call_context_function:
|
|
return keras_call_context_function().layer is not None
|
|
else:
|
|
return False
|
|
|
|
|
|
def output_all_intermediates():
|
|
"""Whether to output all intermediates of a functional control flow op.
|
|
|
|
The default behavior is to output intermediates only when building a Keras
|
|
Layer in graph mode and that too when certain other conditions are met:
|
|
1. We do not output intermediates if the functional control flow op
|
|
is being built inside a FuncGraph which is not a If/While graph. This
|
|
guards against outputting intermediates in eager mode since keras adds
|
|
tensors to a FuncGraph named "keras_graph" in that case. Also because we
|
|
do not output intermediates of tf.function (since this feature is only for
|
|
backwards compatibility) outputting intermediates of functional control
|
|
flow ops built inside tf.function is of no value.
|
|
2. We do not output intermediates when the compilation is using XLA or for a
|
|
TPU.
|
|
3. We do not output intermediates when a single threaded executor is used
|
|
since that does not perform inlining and pruning.
|
|
|
|
Returns:
|
|
A bool telling whether to output all intermediates.
|
|
"""
|
|
if _EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE is not None:
|
|
return _EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE
|
|
if in_defun():
|
|
return False
|
|
if (control_flow_util.GraphOrParentsInXlaContext(ops.get_default_graph()) or
|
|
_is_tpu_strategy(distribution_strategy_context.get_strategy())):
|
|
return False
|
|
if (context.context().function_call_options.executor_type ==
|
|
"SINGLE_THREADED_EXECUTOR"):
|
|
return False
|
|
return _is_building_keras_layer()
|
|
|
|
|
|
def get_func_graph(op, input_shapes, func_name):
|
|
"""Generates and returns a FuncGraph for the given op and input_shapes."""
|
|
fdef = None
|
|
graph = op.graph
|
|
# Recursively search the func in graphs.
|
|
while graph is not None:
|
|
func = graph._get_function(func_name) # pylint: disable=protected-access
|
|
if func is not None:
|
|
fdef = func.definition
|
|
break
|
|
if hasattr(graph, "outer_graph"):
|
|
graph = graph.outer_graph
|
|
else:
|
|
break
|
|
|
|
if fdef is None:
|
|
raise KeyError("%s cannot be found in the graph" % func_name)
|
|
|
|
# `op.graph` may not be the same as `ops.get_default_graph()` e.g.
|
|
# in the case of nested if ops or when the gradient is being computed
|
|
# from inside a Defun. We build the `func_graph` with `op.graph` as its
|
|
# `outer_graph`. This resembles how the `FuncGraph` was built in the
|
|
# forward pass. We need this so that we can resolve references to tensors
|
|
# in `func_graph` from its gradient graph in `_resolve_grad_inputs`.
|
|
with op.graph.as_default():
|
|
func_graph = function_def_to_graph.function_def_to_graph(
|
|
fdef, input_shapes=input_shapes)
|
|
|
|
# TODO(xjun): Ideally we want to retrieve the gradient functions instead of
|
|
# re-create them. But the lifetime of gradient functions of PartitionedCall
|
|
# ops is attached to ParitionedCall ops in the original func_graph and
|
|
# when we are inside this function we don't have access to the original
|
|
# func_graph or PartitionedCall ops. See cl/499362867 and cl/273858076 for
|
|
# more context.
|
|
for operation in func_graph.get_operations():
|
|
if operation.type in ["PartitionedCall", "StatefulPartitionedCall"]:
|
|
f = graph._get_function(operation.get_attr("f").name) # pylint: disable=protected-access
|
|
try:
|
|
cf = function.ConcreteFunction(f.graph, attrs=f.definition.attr)
|
|
except AttributeError:
|
|
# f is not found or f is a _DefinedFunction that doesn't have a graph.
|
|
continue
|
|
operation._gradient_function = cf._get_gradient_function() # pylint: disable=protected-access
|
|
|
|
return func_graph
|
|
|
|
|
|
def get_op_and_outputs(op_or_outputs):
|
|
if isinstance(op_or_outputs, ops.Operation):
|
|
return op_or_outputs, []
|
|
elif not op_or_outputs: # Empty list.
|
|
return None, []
|
|
else:
|
|
return op_or_outputs[0].op, op_or_outputs
|
|
|
|
|
|
def graph_wrapped_for_higher_order_tape_gradients(graph):
|
|
"""Check if `graph` is wrapped by `run_as_function_for_tape_gradients`."""
|
|
while graph is not None:
|
|
if "cflow_gradient_wrapper" in getattr(graph, "name", ""):
|
|
return True
|
|
graph = getattr(graph, "outer_graph", None)
|
|
return False
|
|
|
|
|
|
def run_as_function_for_tape_gradients(make_op, inputs):
|
|
"""Fix higher-order tape gradients by wrapping `make_op` in a function.
|
|
|
|
Args:
|
|
make_op: A function that takes a list of inputs and returns a list of output
|
|
tensors. This function should set any handle data relevant to its outputs
|
|
before returning.
|
|
inputs: A list of tensors to check for tape gradients and pass to
|
|
`make_op`. These should include all tensors used in `make_op`.
|
|
|
|
Returns:
|
|
Tensors corresponding to `make_op`'s output.
|
|
"""
|
|
# GradientTapes created inside a function currently don't work well with
|
|
# un-wrapped control flow ops in that same function. Wrapping in an extra
|
|
# layer of intermediate function means we run extra logic in the function
|
|
# gradient code to record the correct intermediates on the tape.
|
|
#
|
|
# The function attribute inputs to control flow ops are not hashable, so we
|
|
# pass everything as a capture to bypass defun's caching.
|
|
if (gradients_util.PossibleTapeGradientTypes(inputs)
|
|
== gradients_util.POSSIBLE_GRADIENT_TYPES_HIGHER_ORDER
|
|
# We only need one function between the tape and the op; if we've already
|
|
# wrapped once, we stop wrapping to avoid infinite recursion.
|
|
and not (ops.get_default_graph().building_function
|
|
and "cflow_gradient_wrapper" in ops.get_default_graph().name)):
|
|
results = function.defun_with_attributes(
|
|
make_op,
|
|
autograph=False,
|
|
attributes=dict(func_name="cflow_gradient_wrapper"))(inputs)
|
|
return results
|
|
else:
|
|
return make_op(inputs)
|