def testCrossReplicaMean(self):
"""Ensures that cross_replica_mean() executes only when num_shards > 1."""
with ops.Graph().as_default():
with tpu_function.tpu_shard_context(4):
tensor = array_ops.zeros([], dtype=dtypes.float32)
mean = utils.cross_replica_mean(tensor)
self.assertNotEqual(mean, tensor)
with ops.Graph().as_default():
with tpu_function.tpu_shard_context(1):
tensor = array_ops.zeros([], dtype=dtypes.float32)
mean = utils.cross_replica_mean(tensor)
self.assertEqual(mean, tensor)
with ops.Graph().as_default():
with self.assertRaises(ValueError): # Outside of TPU context.
tensor = array_ops.zeros([], dtype=dtypes.float32)
mean = utils.cross_replica_mean(tensor)
def testCrossReplicaMean(self):
"""Ensures that cross_replica_mean() executes only when num_shards > 1."""
with ops.Graph().as_default():
with tpu_function.tpu_shard_context(4):
tensor = array_ops.zeros([], dtype=dtypes.float32)
mean = utils.cross_replica_mean(tensor)
self.assertNotEqual(mean, tensor)
with ops.Graph().as_default():
with tpu_function.tpu_shard_context(1):
tensor = array_ops.zeros([], dtype=dtypes.float32)
mean = utils.cross_replica_mean(tensor)
self.assertEqual(mean, tensor)
with ops.Graph().as_default():
with self.assertRaises(ValueError): # Outside of TPU context.
tensor = array_ops.zeros([], dtype=dtypes.float32)
mean = utils.cross_replica_mean(tensor)
def make_covariance_update_op(self, ema_decay):
"""Constructs and returns the covariance update Op.
Args:
ema_decay: The exponential moving average decay (float or Tensor).
Returns:
An Op for updating the covariance Variable referenced by _cov.
"""
new_cov = math_ops.add_n(
tuple(self._compute_new_cov(idx) for idx in range(self._num_sources)))
# Synchronize value across all TPU cores.
if utils.on_tpu():
new_cov = utils.cross_replica_mean(new_cov)
return moving_averages.assign_moving_average(
self._cov, new_cov, ema_decay, zero_debias=ZERO_DEBIAS)
def make_covariance_update_op(self, ema_decay):
"""Constructs and returns the covariance update Op.
Args:
ema_decay: The exponential moving average decay (float or Tensor).
Returns:
An Op for updating the covariance Variable referenced by _cov.
"""
new_cov = math_ops.add_n(
tuple(
self._compute_new_cov(idx)
for idx in range(self._num_sources)))
# Synchronize value across all TPU cores.
if utils.on_tpu():
new_cov = utils.cross_replica_mean(new_cov)
return moving_averages.assign_moving_average(self._cov,
new_cov,
ema_decay,
zero_debias=ZERO_DEBIAS)
def make_covariance_update_op(self, ema_decay):
"""Constructs and returns the covariance update Op.
Args:
ema_decay: The exponential moving average decay (float or Tensor).
Returns:
An Op for updating the covariance Variable referenced by _cov.
"""
new_cov_contribs = tuple(
self._compute_new_cov(idx) for idx in range(self._num_sources))
# This gets the job done but we might want a better solution in the future.
# In particular, we could have a separate way of specifying where the
# the cov variables finally end up, independent of where their various
# contributions are computed. Right now these are the same thing, but in
# the future we might want to perform the cov computations on each tower,
# so that each tower will be considered a "source" (allowing us to reuse
# the existing "source" code for this).
with maybe_colocate_with(new_cov_contribs[0]):
new_cov = math_ops.add_n(new_cov_contribs)
# Synchronize value across all TPU cores.
if utils.on_tpu():
new_cov = utils.cross_replica_mean(new_cov)
return moving_averages.assign_moving_average(
self._cov, new_cov, ema_decay, zero_debias=ZERO_DEBIAS)
def make_covariance_update_op(self, ema_decay):
"""Constructs and returns the covariance update Op.
Args:
ema_decay: The exponential moving average decay (float or Tensor).
Returns:
An Op for updating the covariance Variable referenced by _cov.
"""
new_cov_contribs = tuple(self._compute_new_cov(idx)
for idx in range(self._num_sources))
# This gets the job done but we might want a better solution in the future.
# In particular, we could have a separate way of specifying where the
# the cov variables finally end up, independent of where their various
# contributions are computed. Right now these are the same thing, but in
# the future we might want to perform the cov computations on each tower,
# so that each tower will be considered a "source" (allowing us to reuse
# the existing "source" code for this).
with maybe_colocate_with(new_cov_contribs[0]):
new_cov = math_ops.add_n(new_cov_contribs)
# Synchronize value across all TPU cores.
if utils.on_tpu():
new_cov = utils.cross_replica_mean(new_cov)
return moving_averages.assign_moving_average(
self._cov, new_cov, ema_decay, zero_debias=ZERO_DEBIAS)