Instance-Level Meta Normalization for Tensorflow 1.1x [https://arxiv.org/abs/1904.03516]

ilm_norm.py

import tensorflow as tf

def group_norm(x, G=32, eps=1e-5):
    # x: input features with shape [N,C,H,W]
    # G: number of groups for GN
    N, H, W, C = x.shape
    gamma = tf.get_variable('gamma', [1, 1, 1, C], initializer=tf.initializers.ones)
    beta = tf.get_variable('B_beta', [1, 1, 1, C], initializer=tf.initializers.zeros)
    x = tf.reshape(x, [N, H, W, G, C // G])
    mean, var = tf.nn.moments(x, [1, 2, 4], keep_dims=True)
    x = (x - mean) / tf.sqrt(var + eps)
    x = tf.reshape(x, [N, C, H, W])
    return x * gamma + beta

def ilm_norm(input, G=16, eps=1e-5):
    with tf.variable_scope(None, default_name='ilm_norm'):
        # Based on: Instance-Level Meta Normalization - https://arxiv.org/abs/1904.03516
        # x: input features with shape [N,H,W,C]
        # G: number of groups for GN
        shape = tf.shape(input)
        N, H, W = shape[0], shape[1], shape[2]
        C = input.shape[-1]
        x = tf.reshape(input, [N, H, W, G, C // G])
        mean, var = tf.nn.moments(x, [1, 2, 4], keep_dims=True)
        x_standartized = (x - mean) / tf.sqrt(var + eps)

        with tf.variable_scope('E', reuse=tf.AUTO_REUSE):
            e_mean = tf.layers.dense(tf.layers.flatten(mean), G, activation=tf.nn.relu)

        with tf.variable_scope('E', reuse=tf.AUTO_REUSE):
            e_var = tf.layers.dense(tf.layers.flatten(var), G, activation=tf.nn.relu)

        with tf.variable_scope('D_mean'):
            d_mean = tf.layers.dense(e_mean, G, activation=tf.nn.tanh)
            d_mean = tf.reshape(d_mean, [N, G, 1, 1])
            d_mean = tf.image.resize_nearest_neighbor(d_mean, (C, 1), align_corners=True)
            d_mean = tf.reshape(d_mean, [N, 1, 1, C])

        with tf.variable_scope('D_var'):
            d_var = tf.layers.dense(e_var, G, activation=tf.nn.sigmoid)
            d_var = tf.reshape(d_var, [N, G, 1, 1])
            d_var = tf.image.resize_nearest_neighbor(d_var, (C, 1), align_corners=True)
            d_var = tf.reshape(d_var, [N, 1, 1, C])

        b_omega = tf.get_variable('B_omega', [1, 1, 1, C], initializer=tf.initializers.ones)
        b_beta = tf.get_variable('B_beta', [1, 1, 1, C], initializer=tf.initializers.zeros)

        omega = d_mean + b_omega
        beta = d_var + b_beta

        return input * omega + beta