# 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)