def fn_with_recompute(*args):
cached_vs.append(variable_scope.get_variable_scope())
# TODO(rsepassi): Rm conditional in TF 1.4
if hasattr(contrib_framework_ops, "current_arg_scope"):
cached_arg_scope.append(contrib_framework_ops.current_arg_scope())
else:
cached_arg_scope.append({})
return fn(*args)
def fn_with_recompute(*args):
cached_vs.append(variable_scope.get_variable_scope())
# TODO(rsepassi): Rm conditional in TF 1.4
if hasattr(contrib_framework_ops, "current_arg_scope"):
cached_arg_scope.append(contrib_framework_ops.current_arg_scope())
else:
cached_arg_scope.append({})
return fn(*args)
def fn_with_recompute(*args):
"""Wrapper for fn."""
# Forward pass
vs = variable_scope.get_variable_scope()
arg_scope = contrib_framework_ops.current_arg_scope()
with backprop.GradientTape() as tape:
fn_kwargs = {}
if has_is_recompute_kwarg:
fn_kwargs["is_recomputing"] = False
outputs = fn(*args, **fn_kwargs)
original_vars = set(tape.watched_variables())
# Backward pass
def grad_fn(*output_grads, **kwargs):
"""Recompute outputs for gradient computation."""
variables = []
if original_vars:
variables = kwargs["variables"]
if set(variables) != original_vars:
raise ValueError(_WRONG_VARS_ERR)
del kwargs
inputs = list(args)
# Recompute outputs
with framework_ops.control_dependencies(output_grads):
if use_data_dep_:
inputs = _force_data_dependency(output_grads, inputs)
with contrib_framework_ops.arg_scope(arg_scope):
with variable_scope.variable_scope(vs, reuse=True):
with backprop.GradientTape() as tape:
fn_kwargs = {}
if has_is_recompute_kwarg:
fn_kwargs["is_recomputing"] = True
outputs = fn(*inputs, **fn_kwargs)
recompute_vars = set(tape.watched_variables())
if original_vars != recompute_vars:
raise ValueError(_WRONG_VARS_ERR)
if not (isinstance(outputs, list) or isinstance(outputs, tuple)):
outputs = [outputs]
outputs = list(outputs)
grads = gradients_impl.gradients(outputs, inputs + variables,
output_grads)
if tupleize_grads:
if use_data_dep_:
grads = _tuple_with_data_dep(grads)
else:
grads = control_flow_ops.tuple(grads)
grad_inputs = grads[:len(inputs)]
grad_vars = grads[len(inputs):]
return grad_inputs, grad_vars
return outputs, grad_fn
def fn_with_recompute(*args):
"""Wrapper for fn."""
# Forward pass
vs = variable_scope.get_variable_scope()
arg_scope = contrib_framework_ops.current_arg_scope()
with backprop.GradientTape() as tape:
fn_kwargs = {}
if has_is_recompute_kwarg:
fn_kwargs["is_recomputing"] = False
outputs = fn(*args, **fn_kwargs)
original_vars = set(tape.watched_variables())
# Backward pass
def grad_fn(*output_grads, **kwargs):
"""Recompute outputs for gradient computation."""
variables = []
if original_vars:
variables = kwargs["variables"]
if set(variables) != original_vars:
raise ValueError(_WRONG_VARS_ERR)
del kwargs
inputs = list(args)
# Recompute outputs
with framework_ops.control_dependencies(output_grads):
if use_data_dep_:
inputs = _force_data_dependency(output_grads, inputs)
with contrib_framework_ops.arg_scope(arg_scope):
with variable_scope.variable_scope(vs, reuse=True):
with backprop.GradientTape() as tape:
fn_kwargs = {}
if has_is_recompute_kwarg:
fn_kwargs["is_recomputing"] = True
outputs = fn(*inputs, **fn_kwargs)
recompute_vars = set(tape.watched_variables())
if original_vars != recompute_vars:
raise ValueError(_WRONG_VARS_ERR)
if not (isinstance(outputs, list) or isinstance(outputs, tuple)):
outputs = [outputs]
outputs = list(outputs)
grads = gradients_impl.gradients(outputs, inputs + variables,
output_grads)
if tupleize_grads:
if use_data_dep_:
grads = _tuple_with_data_dep(grads)
else:
grads = control_flow_ops.tuple(grads)
grad_inputs = grads[:len(inputs)]
grad_vars = grads[len(inputs):]
return grad_inputs, grad_vars
return outputs, grad_fn
def fn_with_recompute(*args):
"""Wrapper for fn."""
# Capture the variable and arg scopes so we can re-enter them when
# recomputing.
vs = variable_scope.get_variable_scope()
arg_scope = contrib_framework_ops.current_arg_scope()
# Track all variables touched in the function.
with backprop.GradientTape() as tape:
fn_kwargs = {}
if has_is_recompute_kwarg:
fn_kwargs["is_recomputing"] = False
outputs = fn(*args, **fn_kwargs)
original_vars = set(tape.watched_variables())
def _grad_fn(output_grads, variables=None):
# Validate that custom_gradient passes the right variables into grad_fn.
if original_vars:
assert variables, (
"Fn created variables but the variables were not "
"passed to the gradient fn.")
if set(variables) != original_vars:
raise ValueError(_WRONG_VARS_ERR)
return _recomputing_grad_fn(
compute_fn=fn,
original_args=args,
original_vars=original_vars,
output_grads=output_grads,
grad_fn_variables=variables,
use_data_dep=use_data_dep_,
tupleize_grads=tupleize_grads,
arg_scope=arg_scope,
var_scope=vs,
has_is_recompute_kwarg=has_is_recompute_kwarg)
# custom_gradient inspects the signature of the function to determine
# whether the user expects variables passed in the grad_fn. If the function
# created variables, the grad_fn should accept the "variables" kwarg.
if original_vars:
def grad_fn(*output_grads, **kwargs):
return _grad_fn(output_grads, kwargs["variables"])
else:
def grad_fn(*output_grads):
return _grad_fn(output_grads)
return outputs, grad_fn
def fn_with_recompute(*args):
"""Wrapper for fn."""
# Capture the variable and arg scopes so we can re-enter them when
# recomputing.
vs = variable_scope.get_variable_scope()
arg_scope = contrib_framework_ops.current_arg_scope()
# Track all variables touched in the function.
with backprop.GradientTape() as tape:
fn_kwargs = {}
if has_is_recompute_kwarg:
fn_kwargs["is_recomputing"] = False
outputs = fn(*args, **fn_kwargs)
original_vars = set(tape.watched_variables())
def _grad_fn(output_grads, variables=None):
# Validate that custom_gradient passes the right variables into grad_fn.
if original_vars:
assert variables, ("Fn created variables but the variables were not "
"passed to the gradient fn.")
if set(variables) != original_vars:
raise ValueError(_WRONG_VARS_ERR)
return _recomputing_grad_fn(
compute_fn=fn,
original_args=args,
original_vars=original_vars,
output_grads=output_grads,
grad_fn_variables=variables,
use_data_dep=use_data_dep_,
tupleize_grads=tupleize_grads,
arg_scope=arg_scope,
var_scope=vs,
has_is_recompute_kwarg=has_is_recompute_kwarg)
# custom_gradient inspects the signature of the function to determine
# whether the user expects variables passed in the grad_fn. If the function
# created variables, the grad_fn should accept the "variables" kwarg.
if original_vars:
def grad_fn(*output_grads, **kwargs):
return _grad_fn(output_grads, kwargs["variables"])
else:
def grad_fn(*output_grads):
return _grad_fn(output_grads)
return outputs, grad_fn
def fn_with_recompute(*args):
"""Wrapper for fn."""
# Forward pass
vs = variable_scope.get_variable_scope()
arg_scope = contrib_framework_ops.current_arg_scope()
with backprop.GradientTape() as tape:
fn_kwargs = {}
if has_is_recompute_kwarg:
fn_kwargs["is_recomputing"] = False
outputs = fn(*args, **fn_kwargs)
original_vars = set(tape.watched_variables())
# Backward pass
def _grad_fn(output_grads, variables=None):
"""Recompute outputs for gradient computation."""
variables = variables or []
if original_vars:
assert variables, ("Fn created variables but the variables were not "
"passed to the gradient fn.")
if set(variables) != original_vars:
raise ValueError(_WRONG_VARS_ERR)
inputs = [array_ops.identity(x) for x in list(args)]
# Recompute outputs
with framework_ops.control_dependencies(output_grads):
if use_data_dep_:
inputs = _force_data_dependency(output_grads, inputs)
with contrib_framework_ops.arg_scope(arg_scope):
with variable_scope.variable_scope(vs, reuse=True):
with backprop.GradientTape() as tape:
fn_kwargs = {}
if has_is_recompute_kwarg:
fn_kwargs["is_recomputing"] = True
outputs = fn(*inputs, **fn_kwargs)
recompute_vars = set(tape.watched_variables())
if original_vars != recompute_vars:
raise ValueError(_WRONG_VARS_ERR)
if not isinstance(outputs, (list, tuple)):
outputs = [outputs]
outputs = list(outputs)
grads = gradients_impl.gradients(outputs, inputs + variables,
output_grads)
if tupleize_grads:
if use_data_dep_:
grads = _tuple_with_data_dep(grads)
else:
grads = control_flow_ops.tuple(grads)
grad_inputs = grads[:len(inputs)]
grad_vars = grads[len(inputs):]
return grad_inputs, grad_vars
# custom_gradient inspects the signature of the function to determine
# whether the user expects variables passed in the grad_fn. If the function
# created variables, the grad_fn should accept the "variables" kwarg.
if original_vars:
def grad_fn(*output_grads, **kwargs):
return _grad_fn(output_grads, kwargs["variables"])
else:
def grad_fn(*output_grads):
return _grad_fn(output_grads)
return outputs, grad_fn
def fn_with_recompute(*args):
"""Wrapper for fn."""
# Forward pass
vs = variable_scope.get_variable_scope()
arg_scope = contrib_framework_ops.current_arg_scope()
with backprop.GradientTape() as tape:
fn_kwargs = {}
if has_is_recompute_kwarg:
fn_kwargs["is_recomputing"] = False
outputs = fn(*args, **fn_kwargs)
original_vars = set(tape.watched_variables())
# Backward pass
def _grad_fn(output_grads, variables=None):
"""Recompute outputs for gradient computation."""
variables = variables or []
if original_vars:
assert variables, (
"Fn created variables but the variables were not "
"passed to the gradient fn.")
if set(variables) != original_vars:
raise ValueError(_WRONG_VARS_ERR)
inputs = [array_ops.identity(x) for x in list(args)]
# Recompute outputs
with framework_ops.control_dependencies(output_grads):
if use_data_dep_:
inputs = _force_data_dependency(output_grads, inputs)
with contrib_framework_ops.arg_scope(arg_scope):
with variable_scope.variable_scope(vs, reuse=True):
with backprop.GradientTape() as tape:
fn_kwargs = {}
if has_is_recompute_kwarg:
fn_kwargs["is_recomputing"] = True
outputs = fn(*inputs, **fn_kwargs)
recompute_vars = set(tape.watched_variables())
if original_vars != recompute_vars:
raise ValueError(_WRONG_VARS_ERR)
if not isinstance(outputs, (list, tuple)):
outputs = [outputs]
outputs = list(outputs)
grads = gradients_impl.gradients(outputs, inputs + variables,
output_grads)
if tupleize_grads:
if use_data_dep_:
grads = _tuple_with_data_dep(grads)
else:
grads = control_flow_ops.tuple(grads)
grad_inputs = grads[:len(inputs)]
grad_vars = grads[len(inputs):]
return grad_inputs, grad_vars
# custom_gradient inspects the signature of the function to determine
# whether the user expects variables passed in the grad_fn. If the function
# created variables, the grad_fn should accept the "variables" kwarg.
if original_vars:
def grad_fn(*output_grads, **kwargs):
return _grad_fn(output_grads, kwargs["variables"])
else:
def grad_fn(*output_grads):
return _grad_fn(output_grads)
return outputs, grad_fn
def fn_with_recompute(*args): cached_vs.append(variable_scope.get_variable_scope()) cached_arg_scope.append(contrib_framework_ops.current_arg_scope()) return fn(*args)