message: Replace the fake updates with no updates when not training. This is possible now that the tf.cond bug has been fixed, and is needed to remove rare data races.
diff: @@ -436,27 +436,30 @@ class BatchNormalization(base.Layer):
  if dmax is not None:
  d = math_ops.maximum(d, -dmax)
  d = math_ops.minimum(d, dmax)
  - # When not training, use r=1, d=0, and decay=1 meaning no updates.
  + # When not training, use r=1, d=0.
  r = utils.smart_cond(training, lambda: r, lambda: array_ops.ones_like(r))
  d = utils.smart_cond(training, lambda: d, lambda: array_ops.zeros_like(d))
  - decay = utils.smart_cond(training, lambda: self.renorm_momentum, lambda: 1.)
  def _update_renorm_variable(var, weight, value):
  """Updates a moving average and weight, returns the unbiased value."""
  + value = array_ops.identity(value)
  + def _do_update():
  # Update the variables without zero debiasing. The debiasing will be
  # accomplished by dividing the exponential moving average by the weight.
  # For example, after a single update, the moving average would be
  # (1-decay) * value. and the weight will be 1-decay, with their ratio
  - # giving value.
  + # giving the value.
  # Make sure the weight is not updated until before r and d computation.
  - value = array_ops.identity(value)
  with ops.control_dependencies([value]):
  weight_value = array_ops.constant(1., dtype=weight.dtype)
  new_var = moving_averages.assign_moving_average(
  - var, value, decay, zero_debias=False)
  + var, value, self.renorm_momentum, zero_debias=False)
  new_weight = moving_averages.assign_moving_average(
  - weight, weight_value, decay, zero_debias=False)
  + weight, weight_value, self.renorm_momentum, zero_debias=False)
  return new_var / new_weight
  + def _fake_update():
  + return array_ops.identity(var)
  + return utils.smart_cond(training, _do_update, _fake_update)
  with ops.colocate_with(self.moving_mean):
  new_mean = _update_renorm_variable(self.renorm_mean,
  @@ -562,8 +565,6 @@ class BatchNormalization(base.Layer):
  else:
  new_mean, new_variance = mean, variance
  - # Update moving averages when training, and prevent updates otherwise.
  - decay = utils.smart_cond(training, lambda: self.momentum, lambda: 1.)
  if self.virtual_batch_size is not None:
  # This isn't strictly correct since in ghost batch norm, you are
  # supposed to sequentially update the moving_mean and moving_variance
  @@ -575,10 +576,18 @@ class BatchNormalization(base.Layer):
  new_variance = math_ops.reduce_mean(new_variance,
  axis=1, keep_dims=True)
  - mean_update = moving_averages.assign_moving_average(
  - self.moving_mean, new_mean, decay, zero_debias=False)
  - variance_update = moving_averages.assign_moving_average(
  - self.moving_variance, new_variance, decay, zero_debias=False)
  + def _do_update(var, value):
  + return moving_averages.assign_moving_average(
  + var, value, self.momentum, zero_debias=False)
  +
  + mean_update = utils.smart_cond(
  + training,
  + lambda: _do_update(self.moving_mean, new_mean),
  + lambda: self.moving_mean)
  + variance_update = utils.smart_cond(
  + training,
  + lambda: _do_update(self.moving_variance, new_variance),
  + lambda: self.moving_variance)
  if context.in_graph_mode():
  self.add_update(mean_update, inputs=inputs)
  self.add_update(variance_update, inputs=inputs)
  