def testCreateSlotFromVariableRespectsScope(self):
# See discussion on #2740.
with self.cached_session():
with variable_scope.variable_scope("scope"):
v = variables.Variable([1.0, 2.5], name="var")
slot = slot_creator.create_slot(v, v.initialized_value(), name="slot")
self.assertEqual("scope/scope/var/slot", slot.op.name)
def _create_lagrangian_multipliers(optimizer_dict, doo_ds):
lag_mul = [slot_creator.create_slot(v, utils.val_or_zero(der, v), 'alpha') for v, der
in zip(optimizer_dict.state, doo_ds)]
[tf.add_to_collection(utils.GraphKeys.LAGRANGIAN_MULTIPLIERS, lm) for lm in lag_mul]
utils.remove_from_collection(utils.GraphKeys.GLOBAL_VARIABLES, *lag_mul)
# this prevents the 'automatic' initialization with tf.global_variables_initializer.
return lag_mul
def apply(self, var_list=None):
# TODO(touts): op_scope
if var_list is None:
var_list = variables.trainable_variables()
for var in var_list:
if var.dtype.base_dtype not in [dtypes.float32, dtypes.float64]:
raise TypeError(
"The variables must be float or double: %s" % var)
if var in self._averages:
raise ValueError(
"Moving average already computed for: %s" % var)
# For variables: to lower communication bandwidth across devices we keep
# the moving averages on the same device as the variables. For other
# tensors, we rely on the existing device allocation mechanism.
if isinstance(var, variables.Variable):
avg = slot_creator.create_slot(
var, var.initialized_value(), self._name,
colocate_with_primary=True)
else:
avg = slot_creator.create_zeros_slot(
var, self._name, colocate_with_primary=(var.op.type == "Variable"))
self._averages[var] = avg
with ops.name_scope(self._name) as scope:
decay = self._num_updates / (self._num_updates + 1)
updates = []
updates.append(self._num_updates_op)
for var in var_list:
updates.append(assign_moving_average(
self._averages[var], var, decay))
return control_flow_ops.group(*updates, name=scope)
def _create_z(optimizer_dict, hyper, d_init_dynamics_d_hyper):
if d_init_dynamics_d_hyper is None: d_init_dynamics_d_hyper = [None]*len(optimizer_dict.state)
with tf.variable_scope('Z'):
z = [slot_creator.create_slot(v, utils.val_or_zero(der, v), hyper.op.name) for v, der
in zip(optimizer_dict.state, d_init_dynamics_d_hyper)]
[tf.add_to_collection(utils.GraphKeys.ZS, lm) for lm in z]
# utils.remove_from_collection(utils.GraphKeys.GLOBAL_VARIABLES, *z)
# in this case it is completely fine to keep zs into the global variable...
return z
def _create_hypergradient(hyper, doo_dhypers):
"""
Creates one hyper-gradient as a variable..
:param hyper: the relative hyperparameter
:param doo_dhypers: initialization, that is the derivative of the outer objective w.r.t this hyper
:return:
"""
hgs = slot_creator.create_slot(hyper, utils.val_or_zero(doo_dhypers, hyper), 'hypergradient')
utils.remove_from_collection(utils.GraphKeys.GLOBAL_VARIABLES, hgs)
return hgs
def testCreateSlotFromVariable(self):
with self.test_session():
v = tf.Variable([1.0, 2.5], name="var")
slot = slot_creator.create_slot(v, v.initialized_value(), name="slot")
tf.initialize_all_variables().run()
self.assertEqual(slot.op.name, "var/slot")
self.assertEqual(slot.get_shape().as_list(), [2])
self.assertEqual(slot.dtype.base_dtype, tf.float32)
self.assertAllEqual(slot.eval(), [1.0, 2.5])
def testCreateSlotFromTensor(self):
with self.cached_session():
v = constant_op.constant([1.0, 2.5], name="const")
slot = slot_creator.create_slot(v, v * 2, name="slot")
variables.global_variables_initializer().run()
self.assertEqual("const/slot", slot.op.name)
self.assertEqual([2], slot.get_shape().as_list())
self.assertEqual(dtypes.float32, slot.dtype.base_dtype)
self.assertAllEqual([2.0, 5.0], self.evaluate(slot))
def testCreateSlotFromTensor(self):
with self.test_session():
v = tf.constant([1.0, 2.5], name="const")
slot = slot_creator.create_slot(v, v * 2, name="slot")
tf.initialize_all_variables().run()
self.assertEqual(slot.op.name, "const/slot")
self.assertEqual(slot.get_shape().as_list(), [2])
self.assertEqual(slot.dtype.base_dtype, tf.float32)
self.assertAllEqual(slot.eval(), [2.0, 5.0])
def testCreateSlotFromVariable(self):
with self.cached_session():
v = variables.Variable([1.0, 2.5], name="var")
slot = slot_creator.create_slot(v, v.initialized_value(), name="slot")
variables.global_variables_initializer().run()
self.assertEqual("var/slot", slot.op.name)
self.assertEqual([2], slot.get_shape().as_list())
self.assertEqual(dtypes.float32, slot.dtype.base_dtype)
self.assertAllEqual([1.0, 2.5], self.evaluate(slot))
def _get_or_make_slot(self, var, val, slot_name, op_name):
"""Find or create a slot for a variable.
Args:
var: A `Variable` object.
val: A `Tensor`. The initial value of the slot.
slot_name: Name for the slot.
op_name: Name to use when scoping the Variable that
needs to be created for the slot.
Returns:
A `Variable` object.
"""
named_slots = self._slot_dict(slot_name)
if var not in named_slots:
named_slots[var] = slot_creator.create_slot(var, val, op_name)
return named_slots[var]
def _process_slot_restoration(self, slot_restoration, variable):
"""Restore a slot variable's value (creating it if necessary)."""
# TODO(allenl): Move this to Optimizer
assert isinstance(self, optimizer_lib.Optimizer)
named_slots = self._slot_dict(slot_restoration.slot_name)
variable_key = optimizer_lib._var_key(variable) # pylint: disable=protected-access
existing_slot_variable = named_slots.get(variable_key, None)
if existing_slot_variable is None:
base_dtype = slot_restoration.value_pointer.dtype.base_dtype
initializer, = io_ops.restore_v2(
prefix=slot_restoration.value_pointer.save_path,
tensor_names=[slot_restoration.value_pointer.checkpoint_key],
shape_and_slices=[""],
dtypes=[base_dtype],
name="checkpoint_initializer")
new_slot_variable = slot_creator.create_slot(variable, initializer,
slot_restoration.slot_name)
if slot_restoration.value_pointer.session is not None:
slot_restoration.value_pointer.session.run(
new_slot_variable.initializer)
named_slots[variable_key] = new_slot_variable
else:
_assign_existing_variable(
existing_slot_variable, value_pointer=slot_restoration.value_pointer)
def apply(self, var_list=None):
"""Maintains moving averages of variables.
`var_list` must be a list of `Variable` or `Tensor` objects. This method
creates shadow variables for all elements of `var_list`. Shadow variables
for `Variable` objects are initialized to the variable's initial value.
For `Tensor` objects, the shadow variables are initialized to 0.
shadow variables are created with `trainable=False` and added to the
`GraphKeys.ALL_VARIABLES` collection. They will be returned by calls to
`tf.all_variables()`.
Returns an op that updates all shadow variables as described above.
Note that `apply()` can be called multiple times with different lists of
variables.
Args:
var_list: A list of Variable or Tensor objects. The variables
and Tensors must be of types float32 or float64.
Returns:
An Operation that updates the moving averages.
Raises:
TypeError: If the arguments are not all float32 or float64.
ValueError: If the moving average of one of the variables is already
being computed.
"""
# TODO(touts): op_scope
if var_list is None:
var_list = variables.trainable_variables()
for var in var_list:
if var.dtype.base_dtype not in [dtypes.float32, dtypes.float64]:
raise TypeError("The variables must be float or double: %s" %
var)
if var in self._averages:
raise ValueError("Moving average already computed for: %s" %
var)
# For variables: to lower communication bandwidth across devices we keep
# the moving averages on the same device as the variables. For other
# tensors, we rely on the existing device allocation mechanism.
with ops.control_dependencies(None):
if isinstance(var, variables.Variable):
avg = slot_creator.create_slot(var,
var.initialized_value(),
self._name,
colocate_with_primary=True)
else:
avg = slot_creator.create_zeros_slot(
var,
self._name,
colocate_with_primary=(var.op.type == "Variable"))
self._averages[var] = avg
with ops.name_scope(self._name) as scope:
decay = ops.convert_to_tensor(self._decay, name="decay")
if self._num_updates is not None:
num_updates = math_ops.cast(self._num_updates,
dtypes.float32,
name="num_updates")
decay = math_ops.minimum(decay, (1.0 + num_updates) /
(10.0 + num_updates))
updates = []
for var in var_list:
updates.append(
assign_moving_average(self._averages[var], var, decay))
return control_flow_ops.group(*updates, name=scope)
def apply(self, var_list=None):
"""Maintains moving averages of variables.
`var_list` must be a list of `Variable` or `Tensor` objects. This method
creates shadow variables for all elements of `var_list`. Shadow variables
for `Variable` objects are initialized to the variable's initial value.
They will be added to the `GraphKeys.MOVING_AVERAGE_VARIABLES` collection.
For `Tensor` objects, the shadow variables are initialized to 0.
shadow variables are created with `trainable=False` and added to the
`GraphKeys.ALL_VARIABLES` collection. They will be returned by calls to
`tf.all_variables()`.
Returns an op that updates all shadow variables as described above.
Note that `apply()` can be called multiple times with different lists of
variables.
Args:
var_list: A list of Variable or Tensor objects. The variables
and Tensors must be of types float32 or float64.
Returns:
An Operation that updates the moving averages.
Raises:
TypeError: If the arguments are not all float32 or float64.
ValueError: If the moving average of one of the variables is already
being computed.
"""
# TODO(touts): op_scope
if var_list is None:
var_list = variables.trainable_variables()
for var in var_list:
if var.dtype.base_dtype not in [dtypes.float32, dtypes.float64]:
raise TypeError("The variables must be float or double: %s" % var.name)
if var in self._averages:
raise ValueError("Moving average already computed for: %s" % var.name)
# For variables: to lower communication bandwidth across devices we keep
# the moving averages on the same device as the variables. For other
# tensors, we rely on the existing device allocation mechanism.
with ops.control_dependencies(None):
if isinstance(var, variables.Variable):
avg = slot_creator.create_slot(
var, var.initialized_value(), self._name,
colocate_with_primary=True)
else:
avg = slot_creator.create_zeros_slot(
var, self._name,
colocate_with_primary=(var.op.type == "Variable"))
self._averages[var] = avg
ops.add_to_collection(ops.GraphKeys.MOVING_AVERAGE_VARIABLES, var)
with ops.name_scope(self._name) as scope:
decay = ops.convert_to_tensor(self._decay, name="decay")
if self._num_updates is not None:
num_updates = math_ops.cast(self._num_updates, dtypes.float32,
name="num_updates")
decay = math_ops.minimum(decay,
(1.0 + num_updates) / (10.0 + num_updates))
updates = []
for var in var_list:
updates.append(assign_moving_average(self._averages[var], var, decay))
return control_flow_ops.group(*updates, name=scope)
def _get_or_make_slot(self, var, val, slot_name, op_name):
named_slots = self._slot_dict(slot_name)
if var not in named_slots:
named_slots[var] = slot_creator.create_slot(var, val, op_name)
return named_slots[var]
def apply(self, var_list=None):
"""Maintains moving averages of variables.
`var_list` must be a list of `Variable` objects. This method
creates shadow variables (holding the moving averages)
for all elements of `var_list`, and
updates the moving averages using the current `var_list` values. Shadow
variables for `Variable` objects are initialized to the variable's initial
value.
Shadow variables are created with `trainable=False`. To access them you
can use the EMA object's `average` method. Note that `EMA` objects are
not trackable by checkpoints, so if you want to checkpoint or restore the
moving variables you will need to manually grab the shadow
variables via `average()` and assign them as `tf.Module` properties or
directly pass them to your `tf.train.Checkpoint`.
Note that `apply()` can be called multiple times. When eager execution is
enabled each call to apply will update the variables once, so this needs to
be called in a loop.
In legacy TF 1.x graphs, this method returns an op that updates all
shadow variables from the current value of their associated variables. In
TF 1.x graphs without automatically control dependencies this op needs to be
manually run.
Args:
var_list: A list of Variable objects. The variables
must be of types bfloat16, float16, float32, or float64.
(In legacy TF 1.x graphs these may be tensors, but this is unsupported
when eager execution is enabled.)
Returns:
An Operation that updates the moving averages.
Raises:
TypeError: If the arguments are not an allowed type.
"""
# TODO(touts): op_scope
if var_list is None:
var_list = variables.trainable_variables()
for v in var_list:
if (isinstance(v, ops.Tensor)
and ops.executing_eagerly_outside_functions()):
raise TypeError(
"tf.train.ExponentialMovingAverage does not support non-Variable"
" tensors when eager execution is enabled.")
zero_debias_true = set() # set of vars to set `zero_debias=True`
for var in var_list:
if var.dtype.base_dtype not in [
dtypes.bfloat16, dtypes.float16, dtypes.float32,
dtypes.float64
]:
raise TypeError(
"The variables must be half, float, or double: %s" %
var.name)
if var.ref() not in self._averages:
# For variables: to lower communication bandwidth across devices we keep
# the moving averages on the same device as the variables. For other
# tensors, we rely on the existing device allocation mechanism.
with ops.init_scope():
if isinstance(var, variables.Variable):
with ops.device(var.device):
initialized_value = var.initialized_value()
avg = slot_creator.create_slot(
var,
initialized_value,
self.name,
colocate_with_primary=True,
copy_xla_sharding=True)
# NOTE(mrry): We only add `tf.Variable` objects to the
# `MOVING_AVERAGE_VARIABLES` collection.
ops.add_to_collection(
ops.GraphKeys.MOVING_AVERAGE_VARIABLES, var)
else:
avg = slot_creator.create_zeros_slot(
var,
self.name,
colocate_with_primary=(var.op.type in [
"Variable", "VariableV2", "VarHandleOp"
]),
copy_xla_sharding=True)
if self._zero_debias:
zero_debias_true.add(avg.ref())
self._averages[var.ref()] = avg
with ops.name_scope(self.name) as scope:
decay = ops.convert_to_tensor(self._decay,
dtype=dtypes.float32,
name="decay")
if self._num_updates is not None:
num_updates = math_ops.cast(self._num_updates,
dtypes.float32,
name="num_updates")
decay = math_ops.minimum(decay, (1.0 + num_updates) /
(10.0 + num_updates))
updates = []
for var in var_list:
avg = self._averages[var.ref()]
zero_debias = avg.ref() in zero_debias_true
updates.append(
assign_moving_average(avg, var, decay, zero_debias))
return control_flow_ops.group(*updates, name=scope)
def _get_or_make_slot(self, var, val, slot_name, op_name):
named_slots = self._slot_dict(slot_name)
if var not in named_slots:
named_slots[var] = slot_creator.create_slot(var, val, op_name)
return named_slots[var]
def apply(self, var_list=None):
"""Maintains moving averages of variables.
`var_list` must be a list of `Variable` or `Tensor` objects. This method
creates shadow variables for all elements of `var_list`. Shadow variables
for `Variable` objects are initialized to the variable's initial value.
They will be added to the `GraphKeys.MOVING_AVERAGE_VARIABLES` collection.
For `Tensor` objects, the shadow variables are initialized to 0 and zero
debiased (see docstring in `assign_moving_average` for more details).
shadow variables are created with `trainable=False` and added to the
`GraphKeys.ALL_VARIABLES` collection. They will be returned by calls to
`tf.global_variables()`.
Returns an op that updates all shadow variables from the current value of
their associated variables.
Note that `apply()` can be called multiple times. When eager execution is
enabled each call to apply will update the variables once, so this needs to
be called in a loop.
Args:
var_list: A list of Variable or Tensor objects. The variables
and Tensors must be of types bfloat16, float16, float32, or float64.
Returns:
An Operation that updates the moving averages.
Raises:
TypeError: If the arguments are not an allowed type.
"""
# TODO(touts): op_scope
if var_list is None:
var_list = variables.trainable_variables()
zero_debias_true = set() # set of vars to set `zero_debias=True`
for var in var_list:
if var.dtype.base_dtype not in [
dtypes.bfloat16, dtypes.float16, dtypes.float32, dtypes.float64
]:
raise TypeError("The variables must be half, float, or double: %s" %
var.name)
if var not in self._averages:
# For variables: to lower communication bandwidth across devices we keep
# the moving averages on the same device as the variables. For other
# tensors, we rely on the existing device allocation mechanism.
with ops.init_scope():
if isinstance(var, variables.Variable):
avg = slot_creator.create_slot(var,
var.initialized_value(),
self.name,
colocate_with_primary=True)
# NOTE(mrry): We only add `tf.Variable` objects to the
# `MOVING_AVERAGE_VARIABLES` collection.
ops.add_to_collection(ops.GraphKeys.MOVING_AVERAGE_VARIABLES, var)
else:
avg = slot_creator.create_zeros_slot(
var,
self.name,
colocate_with_primary=(var.op.type in ["Variable",
"VariableV2",
"VarHandleOp"]))
if self._zero_debias:
zero_debias_true.add(avg)
self._averages[var] = avg
with ops.name_scope(self.name) as scope:
decay = ops.convert_to_tensor(self._decay, name="decay")
if self._num_updates is not None:
num_updates = math_ops.cast(self._num_updates,
dtypes.float32,
name="num_updates")
decay = math_ops.minimum(decay,
(1.0 + num_updates) / (10.0 + num_updates))
updates = []
for var in var_list:
zero_debias = self._averages[var] in zero_debias_true
updates.append(assign_moving_average(
self._averages[var], var, decay, zero_debias=zero_debias))
return control_flow_ops.group(*updates, name=scope)
def apply(self, var_list=None):
"""Maintains moving averages of variables.
`var_list` must be a list of `Variable` or `Tensor` objects. This method
creates shadow variables for all elements of `var_list`. Shadow variables
for `Variable` objects are initialized to the variable's initial value.
They will be added to the `GraphKeys.MOVING_AVERAGE_VARIABLES` collection.
For `Tensor` objects, the shadow variables are initialized to 0 and zero
debiased (see docstring in `assign_moving_average` for more details).
shadow variables are created with `trainable=False` and added to the
`GraphKeys.ALL_VARIABLES` collection. They will be returned by calls to
`tf.global_variables()`.
Returns an op that updates all shadow variables as described above.
Note that `apply()` can be called multiple times with different lists of
variables.
Args:
var_list: A list of Variable or Tensor objects. The variables
and Tensors must be of types float16, float32, or float64.
Returns:
An Operation that updates the moving averages.
Raises:
TypeError: If the arguments are not all float16, float32, or float64.
ValueError: If the moving average of one of the variables is already
being computed.
"""
# TODO(touts): op_scope
if var_list is None:
var_list = variables.trainable_variables()
zero_debias_true = set() # set of vars to set `zero_debias=True`
for var in var_list:
if var.dtype.base_dtype not in [
dtypes.float16, dtypes.float32, dtypes.float64
]:
raise TypeError(
"The variables must be half, float, or double: %s" %
var.name)
if var in self._averages:
raise ValueError("Moving average already computed for: %s" %
var.name)
# For variables: to lower communication bandwidth across devices we keep
# the moving averages on the same device as the variables. For other
# tensors, we rely on the existing device allocation mechanism.
with ops.init_scope():
if isinstance(var, variables.Variable):
avg = slot_creator.create_slot(var,
var.initialized_value(),
self._name,
colocate_with_primary=True)
# NOTE(mrry): We only add `tf.Variable` objects to the
# `MOVING_AVERAGE_VARIABLES` collection.
ops.add_to_collection(
ops.GraphKeys.MOVING_AVERAGE_VARIABLES, var)
else:
avg = slot_creator.create_zeros_slot(
var,
self._name,
colocate_with_primary=(var.op.type in [
"Variable", "VariableV2", "VarHandleOp"
]))
if self._zero_debias:
zero_debias_true.add(avg)
self._averages[var] = avg
with ops.name_scope(self._name) as scope:
decay = ops.convert_to_tensor(self._decay, name="decay")
if self._num_updates is not None:
num_updates = math_ops.cast(self._num_updates,
dtypes.float32,
name="num_updates")
decay = math_ops.minimum(decay, (1.0 + num_updates) /
(10.0 + num_updates))
updates = []
for var in var_list:
zero_debias = self._averages[var] in zero_debias_true
updates.append(
assign_moving_average(self._averages[var],
var,
decay,
zero_debias=zero_debias))
return control_flow_ops.group(*updates, name=scope)
def apply(self, var_list=None):
"""Maintains moving averages of variables.
`var_list` must be a list of `Variable` or `Tensor` objects. This method
creates shadow variables for all elements of `var_list`. Shadow variables
for `Variable` objects are initialized to the variable's initial value.
They will be added to the `GraphKeys.MOVING_AVERAGE_VARIABLES` collection.
For `Tensor` objects, the shadow variables are initialized to 0 and zero
debiased (see docstring in `assign_moving_average` for more details).
shadow variables are created with `trainable=False` and added to the
`GraphKeys.ALL_VARIABLES` collection. They will be returned by calls to
`tf.compat.v1.global_variables()`.
Returns an op that updates all shadow variables from the current value of
their associated variables.
Note that `apply()` can be called multiple times. When eager execution is
enabled each call to apply will update the variables once, so this needs to
be called in a loop.
Args:
var_list: A list of Variable or Tensor objects. The variables and Tensors
must be of types bfloat16, float16, float32, or float64.
Returns:
An Operation that updates the moving averages.
Raises:
TypeError: If the arguments are not an allowed type.
"""
# TODO(touts): op_scope
if var_list is None:
var_list = variables.trainable_variables()
for v in var_list:
if isinstance(v, ops.EagerTensor):
raise TypeError(
"tf.train.ExponentialMovingAverage does not support non-Variable"
" tensors when eager execution is enabled.")
zero_debias_true = set() # set of vars to set `zero_debias=True`
for var in var_list:
if var.dtype.base_dtype not in [
dtypes.bfloat16, dtypes.float16, dtypes.float32,
dtypes.float64
]:
raise TypeError(
"The variables must be half, float, or double: %s" %
var.name)
if var.experimental_ref() not in self._averages:
# For variables: to lower communication bandwidth across devices we keep
# the moving averages on the same device as the variables. For other
# tensors, we rely on the existing device allocation mechanism.
if isinstance(var, variables.Variable):
if ops.executing_eagerly_outside_functions():
init_value = var.read_value()
else:
init_value = var.initialized_value()
avg = slot_creator.create_slot(var,
init_value,
self.name,
colocate_with_primary=True)
# NOTE(mrry): We only add `tf.Variable` objects to the
# `MOVING_AVERAGE_VARIABLES` collection.
ops.add_to_collection(
ops.GraphKeys.MOVING_AVERAGE_VARIABLES, var)
else:
avg = slot_creator.create_zeros_slot(
var,
self.name,
colocate_with_primary=(var.op.type in [
"Variable", "VariableV2", "VarHandleOp"
]))
if self._zero_debias:
zero_debias_true.add(avg.experimental_ref())
self._averages[var.experimental_ref()] = avg
with ops.name_scope(self.name) as scope:
decay = ops.convert_to_tensor(self._decay, name="decay")
if self._num_updates is not None:
num_updates = math_ops.cast(self._num_updates,
dtypes.float32,
name="num_updates")
decay = math_ops.minimum(decay, (1.0 + num_updates) /
(10.0 + num_updates))
updates = []
for var in var_list:
avg = self._averages[var.experimental_ref()]
zero_debias = avg.experimental_ref() in zero_debias_true
updates.append(
assign_moving_average(avg, var, decay, zero_debias))
return control_flow_ops.group(*updates, name=scope)
