Group-Normalization-Tensorflow-测

group normalization with moving average

import tensorflow as tf
import numpy as np

group namalization implementation

def norm(x, norm_type, is_train, G=32, esp=1e-5):
    with tf.variable_scope('{}_norm'.format(norm_type)):
        if norm_type == 'none':
            output = x
        elif norm_type == 'batch':
            output = tf.contrib.layers.batch_norm(
                x, center=True, scale=True, decay=0.999,
                is_training=is_train, updates_collections=None
            )
        elif norm_type == 'group':
            # normalize
            # tranpose: [bs, h, w, c] to [bs, c, h, w] following the paper
            x = tf.transpose(x, [0, 3, 1, 2])
            N, C, H, W = x.get_shape().as_list()
            G = min(G, C)
            x = tf.reshape(x, [N, G, C // G, H, W])
            mean, var = tf.nn.moments(x, [2, 3, 4], keep_dims=True)
            x = (x - mean) / tf.sqrt(var + esp)
            # per channel gamma and beta
            gamma = tf.get_variable('gamma', [C],
                                    initializer=tf.constant_initializer(1.0))
            beta = tf.get_variable('beta', [C],
                                   initializer=tf.constant_initializer(0.0))
            gamma = tf.reshape(gamma, [1, C, 1, 1])
            beta = tf.reshape(beta, [1, C, 1, 1])

            output = tf.reshape(x, [N, C, H, W]) * gamma + beta
            # tranpose: [bs, c, h, w, c] to [bs, h, w, c] following the paper
            output = tf.transpose(output, [0, 2, 3, 1])
        else:
            raise NotImplementedError
    return output

构建数据

input_x = np.arange(180).reshape([2,3,3,10]) # [bs=2, h=3, w=3, c=2]
input_x

array([[[[  0,   1,   2,   3,   4,   5,   6,   7,   8,   9],
         [ 10,  11,  12,  13,  14,  15,  16,  17,  18,  19],
         [ 20,  21,  22,  23,  24,  25,  26,  27,  28,  29]],

        [[ 30,  31,  32,  33,  34,  35,  36,  37,  38,  39],
         [ 40,  41,  42,  43,  44,  45,  46,  47,  48,  49],
         [ 50,  51,  52,  53,  54,  55,  56,  57,  58,  59]],

        [[ 60,  61,  62,  63,  64,  65,  66,  67,  68,  69],
         [ 70,  71,  72,  73,  74,  75,  76,  77,  78,  79],
         [ 80,  81,  82,  83,  84,  85,  86,  87,  88,  89]]],


       [[[ 90,  91,  92,  93,  94,  95,  96,  97,  98,  99],
         [100, 101, 102, 103, 104, 105, 106, 107, 108, 109],
         [110, 111, 112, 113, 114, 115, 116, 117, 118, 119]],

        [[120, 121, 122, 123, 124, 125, 126, 127, 128, 129],
         [130, 131, 132, 133, 134, 135, 136, 137, 138, 139],
         [140, 141, 142, 143, 144, 145, 146, 147, 148, 149]],

        [[150, 151, 152, 153, 154, 155, 156, 157, 158, 159],
         [160, 161, 162, 163, 164, 165, 166, 167, 168, 169],
         [170, 171, 172, 173, 174, 175, 176, 177, 178, 179]]]])

input_x = tf.Variable(input_x,dtype=tf.float32)
input_x

<tf.Variable 'Variable:0' shape=(2, 3, 3, 10) dtype=float32_ref>

tranpose: [bs, h, w, c] to [bs, c, h, w] following the paper

x = tf.transpose(input_x, [0, 3, 1, 2])
N, C, H, W = x.get_shape().as_list()
print(N,C,H,W)

2 10 3 3

G = 5
G = min(G, C)
G

5

x = tf.reshape(x, [N, G, C // G, H, W])
x

<tf.Tensor 'Reshape_1:0' shape=(2, 5, 2, 3, 3) dtype=float32>

mean, var = tf.nn.moments(x, [2, 3, 4], keep_dims=True)
print('mean:',mean)
print('var:',var)

mean: Tensor("moments/mean:0", shape=(2, 5, 1, 1, 1), dtype=float32)
var: Tensor("moments/variance:0", shape=(2, 5, 1, 1, 1), dtype=float32)

esp=1e-5
x = (x - mean) / tf.sqrt(var + esp)
x

<tf.Tensor 'truediv:0' shape=(2, 5, 2, 3, 3) dtype=float32>

per channel gamma and beta

gamma = tf.get_variable('gamma', [C],
                        initializer=tf.constant_initializer(1.0))
beta = tf.get_variable('beta', [C],
                       initializer=tf.constant_initializer(0.0))
gamma = tf.reshape(gamma, [1, C, 1, 1])
beta = tf.reshape(beta, [1, C, 1, 1])
print('gamma:',gamma)
print('beta:',beta)

gamma: Tensor("Reshape_2:0", shape=(1, 10, 1, 1), dtype=float32)
beta: Tensor("Reshape_3:0", shape=(1, 10, 1, 1), dtype=float32)

output = tf.reshape(x, [N, C, H, W]) * gamma + beta
output

<tf.Tensor 'add_1:0' shape=(2, 10, 3, 3) dtype=float32>

tranpose: [bs, c, h, w, c] to [bs, h, w, c] following the paper

output = tf.transpose(output, [0, 2, 3, 1])
output

<tf.Tensor 'transpose_1:0' shape=(2, 3, 3, 10) dtype=float32>