Working TensorFlow 1.15/2 version

[JohnMBrandt]

I've confirmed that this works in Tensorflow 1.15.4 / 2.X As a note, I train remote sensing models on small image patches and deploy them on patches 100x larger. I wasn't able to get good results by training with local context normalization, but found great results by taking a pretrained model with group normalization and dropping LCN in place of GN at deployment time, with a modest (2-3x) increase in window size vs. training patch size

def local_context_norm(x, scope, G=8, esp=1e-5, window_size = 30):
    with tf.variable_scope('{}_norm'.format(scope)):
        # 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()
        n_per_group = C // G
        x_squared = x * x
        integral_img = tf.cumsum(x, axis = 2, exclusive = False)
        integral_img = tf.cumsum(integral_img, axis = 3, exclusive = False)

        integral_img_sq = tf.cumsum(tf.cumsum(x_squared, axis = 2), axis = 3)

        dilation = (1, window_size, window_size)
        integral_img = tf.expand_dims(integral_img, axis = 4)
        integral_img_sq = tf.expand_dims(integral_img_sq, axis = 4)
        kernel = np.array([[[[[1,-1], [-1, 1]]]]])
        c_kernel = np.ones([1, 1, n_per_group, 1, 1])
        sums = tf.layers.Conv3D(
            filters = 1, kernel_size = [1, 2, 2], 
            kernel_initializer = tf.keras.initializers.Constant(kernel),
            strides = [1, 1, 1], dilation_rate = dilation,
            data_format='channels_last', use_bias = False,
            trainable = False,
            )(integral_img)
        sums = tf.layers.Conv3D(
            filters = 1, kernel_size = [n_per_group, 1, 1], 
            kernel_initializer = tf.keras.initializers.Constant(c_kernel),
            strides = [n_per_group, 1, 1],
            data_format='channels_last', use_bias = False,
            trainable = False,
            )(sums)

        squares = tf.layers.Conv3D(
            filters = 1, kernel_size = [1, 2, 2], 
            kernel_initializer = tf.keras.initializers.Constant(kernel),
            strides = [1, 1, 1], dilation_rate = dilation,
            data_format='channels_last', use_bias = False,
            trainable = False,
            )(integral_img_sq)

        squares = tf.layers.Conv3D(
            filters = 1, kernel_size = [n_per_group, 1, 1], 
            kernel_initializer = tf.keras.initializers.Constant(c_kernel),
            strides = [n_per_group, 1, 1], 
            data_format='channels_last', use_bias = False,
            trainable = False,
            )(squares)
   
        sums = tf.transpose(sums, [0, 4, 1, 2, 3])
        squares = tf.transpose(squares, [0, 4, 1, 2, 3])

        n = window_size * window_size * n_per_group
        means = tf.squeeze(sums / n, axis = 1)
        var = tf.squeeze((1.0 / n * (squares - sums * sums / n)), axis = 1)
        print(means.shape, var.shape)
        _, _, h, w = means.get_shape().as_list()

        pad2d = (int(math.floor((W - w) / 2)), int(math.ceil((W - w) / 2)), int(math.floor((H - h) / 2)),
                     int(math.ceil((H - h) / 2)))
        padded_means = tf.pad(means,[[0,0], [0,0], [pad2d[0], pad2d[1]], [pad2d[2], pad2d[3]]], 'REFLECT')
        padded_vars = tf.pad(var, [[0,0], [0,0], [pad2d[0], pad2d[1]], [pad2d[2], pad2d[3]]], 'REFLECT')

        G = min(G, C)
        x = tf.reshape(x, [-1, G, C // G, H, W])
        padded_means = tf.repeat(padded_means, axis = 1, repeats = n_per_group)
        padded_means = tf.reshape(padded_means, [-1, G, C // G, H, W])
        
        padded_vars = tf.repeat(padded_vars, axis = 1, repeats = n_per_group)
        padded_vars = tf.reshape(padded_vars, [-1, G, C // G, H, W])

        x = (x - padded_means) / tf.sqrt(padded_vars + esp)
        # per channel gamma and beta
        zeros = lambda: tf.zeros([C], dtype=tf.float32)
        ones = lambda: tf.ones([C], dtype=tf.float32)
        gamma = tf.Variable(initial_value = ones, dtype=tf.float32, name=f'gamma_{scope}')
        beta = tf.Variable(initial_value = zeros, dtype=tf.float32, name=f'beta_{scope}')
        gamma = tf.reshape(gamma, [1, C, 1, 1])
        beta = tf.reshape(beta, [1, C, 1, 1])

        output = tf.reshape(x, [-1, 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])
    return output