path: root/lib/virtual_batchnorm_impl.py

# Copyright 2017 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.
# ==============================================================================
"""Virtual batch normalization.

This technique was first introduced in `Improved Techniques for Training GANs`
(Salimans et al, https://arxiv.org/abs/1606.03498). Instead of using batch
normalization on a minibatch, it fixes a reference subset of the data to use for
calculating normalization statistics.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import variable_scope

__all__ = [
    'VBN',
]


def _static_or_dynamic_batch_size(tensor, batch_axis):
  """Returns the static or dynamic batch size."""
  batch_size = array_ops.shape(tensor)[batch_axis]
  static_batch_size = tensor_util.constant_value(batch_size)
  return static_batch_size or batch_size


def _statistics(x, axes):
  """Calculate the mean and mean square of `x`.

  Modified from the implementation of `tf.nn.moments`.

  Args:
    x: A `Tensor`.
    axes: Array of ints.  Axes along which to compute mean and
      variance.

  Returns:
    Two `Tensor` objects: `mean` and `square mean`.
  """
  # The dynamic range of fp16 is too limited to support the collection of
  # sufficient statistics. As a workaround we simply perform the operations
  # on 32-bit floats before converting the mean and variance back to fp16
  y = math_ops.cast(x, dtypes.float32) if x.dtype == dtypes.float16 else x

  # Compute true mean while keeping the dims for proper broadcasting.
  shift = array_ops.stop_gradient(math_ops.reduce_mean(y, axes, keepdims=True))

  shifted_mean = math_ops.reduce_mean(y - shift, axes, keepdims=True)
  mean = shifted_mean + shift
  mean_squared = math_ops.reduce_mean(math_ops.square(y), axes, keepdims=True)

  mean = array_ops.squeeze(mean, axes)
  mean_squared = array_ops.squeeze(mean_squared, axes)
  if x.dtype == dtypes.float16:
    return (math_ops.cast(mean, dtypes.float16),
            math_ops.cast(mean_squared, dtypes.float16))
  else:
    return (mean, mean_squared)


def _validate_init_input_and_get_axis(reference_batch, axis):
  """Validate input and return the used axis value."""
  if reference_batch.shape.ndims is None:
    raise ValueError('`reference_batch` has unknown dimensions.')

  ndims = reference_batch.shape.ndims
  if axis < 0:
    used_axis = ndims + axis
  else:
    used_axis = axis
  if used_axis < 0 or used_axis >= ndims:
    raise ValueError('Value of `axis` argument ' + str(used_axis) +
                     ' is out of range for input with rank ' + str(ndims))
  return used_axis


def _validate_call_input(tensor_list, batch_dim):
  """Verifies that tensor shapes are compatible, except for `batch_dim`."""
  def _get_shape(tensor):
    shape = tensor.shape.as_list()
    del shape[batch_dim]
    return shape
  base_shape = tensor_shape.TensorShape(_get_shape(tensor_list[0]))
  for tensor in tensor_list:
    base_shape.assert_is_compatible_with(_get_shape(tensor))


class VBN(object):
  """A class to perform virtual batch normalization.

  This technique was first introduced in `Improved Techniques for Training GANs`
  (Salimans et al, https://arxiv.org/abs/1606.03498). Instead of using batch
  normalization on a minibatch, it fixes a reference subset of the data to use
  for calculating normalization statistics.

  To do this, we calculate the reference batch mean and mean square, and modify
  those statistics for each example. We use mean square instead of variance,
  since it is linear.

  Note that if `center` or `scale` variables are created, they are shared
  between all calls to this object.

  The `__init__` API is intended to mimic `tf.layers.batch_normalization` as
  closely as possible.
  """

  def __init__(self,
               reference_batch,
               axis=-1,
               epsilon=1e-3,
               center=True,
               scale=True,
               beta_initializer=init_ops.zeros_initializer(),
               gamma_initializer=init_ops.ones_initializer(),
               beta_regularizer=None,
               gamma_regularizer=None,
               trainable=True,
               name=None,
               batch_axis=0):
    """Initialize virtual batch normalization object.

    We precompute the 'mean' and 'mean squared' of the reference batch, so that
    `__call__` is efficient. This means that the axis must be supplied when the
    object is created, not when it is called.

    We precompute 'square mean' instead of 'variance', because the square mean
    can be easily adjusted on a per-example basis.

    Args:
      reference_batch: A minibatch tensors. This will form the reference data
        from which the normalization statistics are calculated. See
        https://arxiv.org/abs/1606.03498 for more details.
      axis: Integer, the axis that should be normalized (typically the features
        axis). For instance, after a `Convolution2D` layer with
        `data_format="channels_first"`, set `axis=1` in `BatchNormalization`.
      epsilon: Small float added to variance to avoid dividing by zero.
      center: If True, add offset of `beta` to normalized tensor. If False,
        `beta` is ignored.
      scale: If True, multiply by `gamma`. If False, `gamma` is
        not used. When the next layer is linear (also e.g. `nn.relu`), this can
        be disabled since the scaling can be done by the next layer.
      beta_initializer: Initializer for the beta weight.
      gamma_initializer: Initializer for the gamma weight.
      beta_regularizer: Optional regularizer for the beta weight.
      gamma_regularizer: Optional regularizer for the gamma weight.
      trainable: Boolean, if `True` also add variables to the graph collection
        `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
      name: String, the name of the ops.
      batch_axis: The axis of the batch dimension. This dimension is treated
        differently in `virtual batch normalization` vs `batch normalization`.

    Raises:
      ValueError: If `reference_batch` has unknown dimensions at graph
        construction.
      ValueError: If `batch_axis` is the same as `axis`.
    """
    axis = _validate_init_input_and_get_axis(reference_batch, axis)
    self._epsilon = epsilon
    self._beta = 0
    self._gamma = 1
    self._batch_axis = _validate_init_input_and_get_axis(
        reference_batch, batch_axis)

    if axis == self._batch_axis:
      raise ValueError('`axis` and `batch_axis` cannot be the same.')

    with variable_scope.variable_scope(name, 'VBN',
                                       values=[reference_batch]) as self._vs:
      self._reference_batch = reference_batch

      # Calculate important shapes:
      #  1) Reduction axes for the reference batch
      #  2) Broadcast shape, if necessary
      #  3) Reduction axes for the virtual batchnormed batch
      #  4) Shape for optional parameters
      input_shape = self._reference_batch.shape
      ndims = input_shape.ndims
      reduction_axes = list(range(ndims))
      del reduction_axes[axis]

      self._broadcast_shape = [1] * len(input_shape)
      self._broadcast_shape[axis] = input_shape[axis].value

      self._example_reduction_axes = list(range(ndims))
      del self._example_reduction_axes[max(axis, self._batch_axis)]
      del self._example_reduction_axes[min(axis, self._batch_axis)]

      params_shape = self._reference_batch.shape[axis]

      # Determines whether broadcasting is needed. This is slightly different
      # than in the `nn.batch_normalization` case, due to `batch_dim`.
      self._needs_broadcasting = (
          sorted(self._example_reduction_axes) != list(range(ndims))[:-2])

      # Calculate the sufficient statistics for the reference batch in a way
      # that can be easily modified by additional examples.
      self._ref_mean, self._ref_mean_squares = _statistics(
          self._reference_batch, reduction_axes)
      self._ref_variance = (self._ref_mean_squares -
                            math_ops.square(self._ref_mean))

      # Virtual batch normalization uses a weighted average between example
      # statistics and the reference batch statistics.
      ref_batch_size = _static_or_dynamic_batch_size(
          self._reference_batch, self._batch_axis)
      self._example_weight = 1. / (math_ops.to_float(ref_batch_size) + 1.)
      self._ref_weight = 1. - self._example_weight

      # Make the variables, if necessary.
      if center:
        self._beta = variable_scope.get_variable(
            name='beta',
            shape=(params_shape,),
            initializer=beta_initializer,
            regularizer=beta_regularizer,
            trainable=trainable)
      if scale:
        self._gamma = variable_scope.get_variable(
            name='gamma',
            shape=(params_shape,),
            initializer=gamma_initializer,
            regularizer=gamma_regularizer,
            trainable=trainable)

  def _virtual_statistics(self, inputs, reduction_axes):
    """Compute the statistics needed for virtual batch normalization."""
    cur_mean, cur_mean_sq = _statistics(inputs, reduction_axes)
    vb_mean = (self._example_weight * cur_mean +
               self._ref_weight * self._ref_mean)
    vb_mean_sq = (self._example_weight * cur_mean_sq +
                  self._ref_weight * self._ref_mean_squares)
    return (vb_mean, vb_mean_sq)

  def _broadcast(self, v, broadcast_shape=None):
    # The exact broadcast shape depends on the current batch, not the reference
    # batch, unless we're calculating the batch normalization of the reference
    # batch.
    b_shape = broadcast_shape or self._broadcast_shape
    if self._needs_broadcasting and v is not None:
      return array_ops.reshape(v, b_shape)
    return v

  def reference_batch_normalization(self):
    """Return the reference batch, but batch normalized."""
    with ops.name_scope(self._vs.name):
      return nn.batch_normalization(self._reference_batch,
                                    self._broadcast(self._ref_mean),
                                    self._broadcast(self._ref_variance),
                                    self._broadcast(self._beta),
                                    self._broadcast(self._gamma),
                                    self._epsilon)

  def __call__(self, inputs):
    """Run virtual batch normalization on inputs.

    Args:
      inputs: Tensor input.

    Returns:
       A virtual batch normalized version of `inputs`.

    Raises:
       ValueError: If `inputs` shape isn't compatible with the reference batch.
    """
    _validate_call_input([inputs, self._reference_batch], self._batch_axis)

    with ops.name_scope(self._vs.name, values=[inputs, self._reference_batch]):
      # Calculate the statistics on the current input on a per-example basis.
      vb_mean, vb_mean_sq = self._virtual_statistics(
          inputs, self._example_reduction_axes)
      vb_variance = vb_mean_sq - math_ops.square(vb_mean)

      # The exact broadcast shape of the input statistic Tensors depends on the
      # current batch, not the reference batch. The parameter broadcast shape
      # is independent of the shape of the input statistic Tensor dimensions.
      b_shape = self._broadcast_shape[:]  # deep copy
      b_shape[self._batch_axis] = _static_or_dynamic_batch_size(
          inputs, self._batch_axis)
      return nn.batch_normalization(
          inputs,
          self._broadcast(vb_mean, b_shape),
          self._broadcast(vb_variance, b_shape),
          self._broadcast(self._beta, self._broadcast_shape),
          self._broadcast(self._gamma, self._broadcast_shape),
          self._epsilon)