evonorm.py

import tensorflow as tf

EPSILON = 1e-5

def _get_shape_list(tensor):
  """Returns tensor's shape as a list which can be unpacked."""
  static_shape = tensor.shape.as_list()
  if not any([x is None for x in static_shape]):
    return static_shape

  dynamic_shape = tf.shape(tensor)
  ndims = tensor.shape.ndims

  # Return mixture of static and dynamic dims.
  shapes = [
      static_shape[i] if static_shape[i] is not None else dynamic_shape[i]
      for i in range(ndims)
  ]
  return shapes

def _group_std(inputs,
               epsilon=EPSILON,
               data_format='channels_first',
               num_groups=32):
  """Grouped standard deviation along the channel dimension."""
  axis = 3 if data_format == 'channels_last' else 1
  while num_groups > 1:
    if inputs.shape[axis] % num_groups == 0:
      break
    num_groups -= 1
  if data_format == 'channels_last':
    _, h, w, c = inputs.shape.as_list()
    x = tf.reshape(inputs, [-1, h, w, num_groups, c // num_groups])
    _, variance = tf.nn.moments(x, [1, 2, 4], keepdims=True)
  else:
    _, c, h, w = inputs.shape.as_list()
    x = tf.reshape(inputs, [-1, num_groups, c // num_groups, h, w])
    _, variance = tf.nn.moments(x, [2, 3, 4], keepdims=True)
  std = tf.sqrt(variance + epsilon)
  std = tf.broadcast_to(std, _get_shape_list(x))
  return tf.reshape(std, _get_shape_list(inputs))

# def evonorm(inputs,
#             layer=LAYER_EVONORM_B0,
#             nonlinearity=True,
#             init_zero=False,
#             epsilon=EPSILON,
#             num_groups=32,
#             data_format='channels_first'):
#   """Apply an EvoNorm transformation (an alternative to BN-ReLU).
#      Hanxiao Liu, Andrew Brock, Karen Simonyan, Quoc V. Le.
#      Evolving Normalization-Activation Layers.
#      https://arxiv.org/abs/2004.02967
#   Args:
#     inputs: `Tensor` whose shape is either `[batch, channels, ...]` with
#         the "channels_first" format or `[batch, height, width, channels]`
#         with the "channels_last" format.
#     is_training: `bool` for whether the model is training.
#     layer: `String` specifies the EvoNorm instantiation.
#     nonlinearity: `bool` if False, apply an affine transform only.
#     init_zero: `bool` if True, initializes scale parameter of batch
#         normalization with 0 instead of 1 (default).
#     decay: `float` a scalar decay used in the moving average.
#     epsilon: `float` a small float added to variance to avoid dividing by zero.
#     num_groups: `int` the number of groups per layer, used only when `layer` ==
#         LAYER_EVONORM_S0.
#     data_format: `str` either "channels_first" for `[batch, channels, height,
#         width]` or "channels_last for `[batch, height, width, channels]`.
#   Returns:
#     A normalized `Tensor` with the same `data_format`.
#   """
#   if init_zero:
#     gamma_initializer = tf.zeros_initializer()
#   else:
#     gamma_initializer = tf.ones_initializer()

#   if data_format == 'channels_last':
#     var_shape = (1, 1, 1, inputs.shape[3])
#   else:
#     var_shape = (1, inputs.shape[1], 1, 1)
#   with tf.variable_scope(None, default_name='evonorm'):
#     beta = tf.get_variable(
#         'beta',
#         shape=var_shape,
#         dtype=inputs.dtype,
#         initializer=tf.zeros_initializer())
#     gamma = tf.get_variable(
#         'gamma',
#         shape=var_shape,
#         dtype=inputs.dtype,
#         initializer=gamma_initializer)
#     if nonlinearity:
#       v = tf.get_variable(
#           'v',
#           shape=var_shape,
#           dtype=inputs.dtype,
#           initializer=tf.ones_initializer())
#       if layer == LAYER_EVONORM_S0:
#         den = _group_std(
#             inputs,
#             epsilon=epsilon,
#             data_format=data_format,
#             num_groups=num_groups)
#         inputs = inputs * tf.nn.sigmoid(v * inputs) / den
#       else:
#         raise ValueError('Unknown EvoNorm layer: {}'.format(layer))
#   return inputs * gamma + beta

class EvoNormS0(tf.keras.layers.Layer):
    def __init__(self, 
    nonlinearity=True, 
    init_zero=False, 
    epsilon=EPSILON, 
    num_groups=32, 
    data_format='channels_first'):
        self.nonlinearity = nonlinearity
        self.init_zero = init_zero
        self.epsilon = epsilon
        self.num_groups = num_groups
        self.data_format = data_format

        super(EvoNormS0, self).__init__()

    def build(self, input_shape):
        if self.init_zero:
            gamma_initializer = tf.zeros_initializer()
        else:
            gamma_initializer = tf.ones_initializer()

        if self.data_format == 'channels_last':
            var_shape = (1, 1, 1, input_shape[3])
        else:
            var_shape = (1, input_shape[1], 1, 1)

        self.beta = tf.Variable(
          name='beta',
          initial_value=tf.zeros_initializer()(var_shape))
        self.gamma = tf.Variable(
          name='gamma',
          initial_value=gamma_initializer(var_shape))
        if self.nonlinearity:
          self.v = tf.Variable(
            name='v',
            initial_value=tf.ones_initializer()(var_shape))


    def call(self, inputs):
      if self.nonlinearity:
        den = _group_std(
            inputs,
            epsilon=self.epsilon,
            data_format=self.data_format,
            num_groups=self.num_groups)
        inputs = inputs * tf.nn.sigmoid(self.v * inputs) / den

      return inputs * self.gamma + self.beta