From 8ef69698b851a4b115227a7e5a079d4577f2ee29 Mon Sep 17 00:00:00 2001 From: Somshubra Majumdar Date: Sat, 15 Apr 2017 11:43:55 -0500 Subject: [PATCH] Update BatchRenorm to Keras 2 API (#69) * Update BatchRenorm to Keras 2 API * update tests * Remove mode 1 test (mode 1 doesnt exist anymore) --- keras_contrib/layers/normalization.py | 251 +++++++++--------- .../layers/test_normalization.py | 34 +-- 2 files changed, 136 insertions(+), 149 deletions(-) diff --git a/keras_contrib/layers/normalization.py b/keras_contrib/layers/normalization.py index eda985a..20eaee5 100644 --- a/keras_contrib/layers/normalization.py +++ b/keras_contrib/layers/normalization.py @@ -1,5 +1,5 @@ from keras.engine import Layer, InputSpec -from .. import initializers, regularizers +from .. import initializers, regularizers, constraints from .. import backend as K from keras.utils.generic_utils import get_custom_objects @@ -14,29 +14,20 @@ class BatchRenormalization(Layer): close to 0 and the activation standard deviation close to 1. # Arguments - epsilon: small float > 0. Fuzz parameter. - Theano expects epsilon >= 1e-5. - mode: integer, 0, 1 or 2. - - 0: feature-wise normalization. - Each feature map in the input will - be normalized separately. The axis on which - to normalize is specified by the `axis` argument. - Note that if the input is a 4D image tensor - using Theano conventions (samples, channels, rows, cols) - then you should set `axis` to `1` to normalize along - the channels axis. - During training and testing we use running averages - computed during the training phase to normalize the data - - 1: sample-wise normalization. This mode assumes a 2D input. - - 2: feature-wise normalization, like mode 0, but - using per-batch statistics to normalize the data during both - testing and training. - axis: integer, axis along which to normalize in mode 0. For instance, - if your input tensor has shape (samples, channels, rows, cols), - set axis to 1 to normalize per feature map (channels axis). + axis: Integer, the axis that should be normalized + (typically the features axis). + For instance, after a `Conv2D` layer with + `data_format="channels_first"`, + set `axis=1` in `BatchRenormalization`. momentum: momentum in the computation of the exponential average of the mean and standard deviation of the data, for feature-wise normalization. + 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. + epsilon: small float > 0. Fuzz parameter. + Theano expects epsilon >= 1e-5. r_max_value: Upper limit of the value of r_max. d_max_value: Upper limit of the value of d_max. t_delta: At each iteration, increment the value of t by t_delta. @@ -44,18 +35,22 @@ class BatchRenormalization(Layer): List of 2 Numpy arrays, with shapes: `[(input_shape,), (input_shape,)]` Note that the order of this list is [gamma, beta, mean, std] - beta_init: name of initialization function for shift parameter + beta_initializer: name of initialization function for shift parameter (see [initializers](../initializers.md)), or alternatively, Theano/TensorFlow function to use for weights initialization. This parameter is only relevant if you don't pass a `weights` argument. - gamma_init: name of initialization function for scale parameter (see + gamma_initializer: name of initialization function for scale parameter (see [initializers](../initializers.md)), or alternatively, Theano/TensorFlow function to use for weights initialization. This parameter is only relevant if you don't pass a `weights` argument. + moving_mean_initializer: Initializer for the moving mean. + moving_variance_initializer: Initializer for the moving variance. gamma_regularizer: instance of [WeightRegularizer](../regularizers.md) (eg. L1 or L2 regularization), applied to the gamma vector. beta_regularizer: instance of [WeightRegularizer](../regularizers.md), applied to the beta vector. + beta_constraint: Optional constraint for the beta weight. + gamma_constraint: Optional constraint for the gamma weight. # Input shape Arbitrary. Use the keyword argument `input_shape` @@ -69,16 +64,16 @@ class BatchRenormalization(Layer): - [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://arxiv.org/abs/1502.03167) """ - def __init__(self, epsilon=1e-3, mode=0, axis=-1, momentum=0.99, - r_max_value=3., d_max_value=5., t_delta=1., weights=None, beta_init='zero', - gamma_init='one', gamma_regularizer=None, beta_regularizer=None, - **kwargs): + def __init__(self, axis=-1, momentum=0.99, center=True, scale=True, epsilon=1e-3, + r_max_value=3., d_max_value=5., t_delta=1., weights=None, beta_initializer='zero', + gamma_initializer='one', moving_mean_initializer='zeros', + moving_variance_initializer='ones', gamma_regularizer=None, beta_regularizer=None, + beta_constraint=None, gamma_constraint=None, **kwargs): self.supports_masking = True - self.beta_init = initializers.get(beta_init) - self.gamma_init = initializers.get(gamma_init) - self.epsilon = epsilon - self.mode = mode self.axis = axis + self.epsilon = epsilon + self.center = center + self.scale = scale self.momentum = momentum self.gamma_regularizer = regularizers.get(gamma_regularizer) self.beta_regularizer = regularizers.get(beta_regularizer) @@ -86,29 +81,51 @@ class BatchRenormalization(Layer): self.r_max_value = r_max_value self.d_max_value = d_max_value self.t_delta = t_delta - if self.mode == 0: - self.uses_learning_phase = True + self.beta_initializer = initializers.get(beta_initializer) + self.gamma_initializer = initializers.get(gamma_initializer) + self.moving_mean_initializer = initializers.get(moving_mean_initializer) + self.moving_variance_initializer = initializers.get(moving_variance_initializer) + self.beta_constraint = constraints.get(beta_constraint) + self.gamma_constraint = constraints.get(gamma_constraint) + super(BatchRenormalization, self).__init__(**kwargs) def build(self, input_shape): - self.input_spec = [InputSpec(shape=input_shape)] - shape = (input_shape[self.axis],) + dim = input_shape[self.axis] + if dim is None: + raise ValueError('Axis ' + str(self.axis) + ' of ' + 'input tensor should have a defined dimension ' + 'but the layer received an input with shape ' + + str(input_shape) + '.') + self.input_spec = InputSpec(ndim=len(input_shape), + axes={self.axis: dim}) + shape = (dim,) - self.gamma = self.add_weight(shape, - initializer=self.gamma_init, - regularizer=self.gamma_regularizer, - name='{}_gamma'.format(self.name)) - self.beta = self.add_weight(shape, - initializer=self.beta_init, - regularizer=self.beta_regularizer, - name='{}_beta'.format(self.name)) - self.running_mean = self.add_weight(shape, initializer='zero', + if self.scale: + self.gamma = self.add_weight(shape, + initializer=self.gamma_initializer, + regularizer=self.gamma_regularizer, + constraint=self.gamma_constraint, + name='{}_gamma'.format(self.name)) + else: + self.gamma = None + + if self.center: + self.beta = self.add_weight(shape, + initializer=self.beta_initializer, + regularizer=self.beta_regularizer, + constraint=self.beta_constraint, + name='{}_beta'.format(self.name)) + else: + self.beta = None + + self.running_mean = self.add_weight(shape, initializer=self.moving_mean_initializer, name='{}_running_mean'.format(self.name), trainable=False) - # Note: running_std actually holds the running variance, not the running std. - self.running_std = self.add_weight(shape, initializer='one', - name='{}_running_std'.format(self.name), - trainable=False) + + self.running_variance = self.add_weight(shape, initializer=self.moving_variance_initializer, + name='{}_running_std'.format(self.name), + trainable=False) self.r_max = K.variable(np.ones((1,)), name='{}_r_max'.format(self.name)) @@ -119,113 +136,97 @@ class BatchRenormalization(Layer): if self.initial_weights is not None: self.set_weights(self.initial_weights) del self.initial_weights + self.built = True - def call(self, x, mask=None): - if self.mode == 0 or self.mode == 2: - assert self.built, 'Layer must be built before being called' - input_shape = K.int_shape(x) + def call(self, inputs, training=None): + assert self.built, 'Layer must be built before being called' + input_shape = K.int_shape(inputs) - reduction_axes = list(range(len(input_shape))) - del reduction_axes[self.axis] - broadcast_shape = [1] * len(input_shape) - broadcast_shape[self.axis] = input_shape[self.axis] + reduction_axes = list(range(len(input_shape))) + del reduction_axes[self.axis] + broadcast_shape = [1] * len(input_shape) + broadcast_shape[self.axis] = input_shape[self.axis] - mean_batch, var_batch = K.moments(x, reduction_axes, shift=None, keep_dims=False) - std_batch = (K.sqrt(var_batch + self.epsilon)) + mean_batch, var_batch = K.moments(inputs, reduction_axes, shift=None, keep_dims=False) + std_batch = (K.sqrt(var_batch + self.epsilon)) - r_max_value = K.get_value(self.r_max) - r = std_batch / (K.sqrt(self.running_std + self.epsilon)) - r = K.stop_gradient(K.clip(r, 1 / r_max_value, r_max_value)) + r_max_value = K.get_value(self.r_max) + r = std_batch / (K.sqrt(self.running_variance + self.epsilon)) + r = K.stop_gradient(K.clip(r, 1 / r_max_value, r_max_value)) - d_max_value = K.get_value(self.d_max) - d = (mean_batch - self.running_mean) / K.sqrt(self.running_std + self.epsilon) - d = K.stop_gradient(K.clip(d, -d_max_value, d_max_value)) + d_max_value = K.get_value(self.d_max) + d = (mean_batch - self.running_mean) / K.sqrt(self.running_variance + self.epsilon) + d = K.stop_gradient(K.clip(d, -d_max_value, d_max_value)) - if sorted(reduction_axes) == range(K.ndim(x))[:-1]: - x_normed_batch = (x - mean_batch) / std_batch - x_normed = (x_normed_batch * r + d) * self.gamma + self.beta - else: - # need broadcasting - broadcast_mean = K.reshape(mean_batch, broadcast_shape) - broadcast_std = K.reshape(std_batch, broadcast_shape) - broadcast_r = K.reshape(r, broadcast_shape) - broadcast_d = K.reshape(d, broadcast_shape) - broadcast_beta = K.reshape(self.beta, broadcast_shape) - broadcast_gamma = K.reshape(self.gamma, broadcast_shape) + if sorted(reduction_axes) == range(K.ndim(inputs))[:-1]: + x_normed_batch = (inputs - mean_batch) / std_batch + x_normed = (x_normed_batch * r + d) * self.gamma + self.beta + else: + # need broadcasting + broadcast_mean = K.reshape(mean_batch, broadcast_shape) + broadcast_std = K.reshape(std_batch, broadcast_shape) + broadcast_r = K.reshape(r, broadcast_shape) + broadcast_d = K.reshape(d, broadcast_shape) + broadcast_beta = K.reshape(self.beta, broadcast_shape) + broadcast_gamma = K.reshape(self.gamma, broadcast_shape) - x_normed_batch = (x - broadcast_mean) / broadcast_std - x_normed = (x_normed_batch * broadcast_r + broadcast_d) * broadcast_gamma + broadcast_beta + x_normed_batch = (inputs - broadcast_mean) / broadcast_std + x_normed = (x_normed_batch * broadcast_r + broadcast_d) * broadcast_gamma + broadcast_beta - # explicit update to moving mean and standard deviation - self.add_update([K.moving_average_update(self.running_mean, mean_batch, self.momentum), - K.moving_average_update(self.running_std, std_batch ** 2, self.momentum)], x) + # explicit update to moving mean and standard deviation + self.add_update([K.moving_average_update(self.running_mean, mean_batch, self.momentum), + K.moving_average_update(self.running_variance, std_batch ** 2, self.momentum)], inputs) - # update r_max and d_max - t_val = K.get_value(self.t) - r_val = self.r_max_value / (1 + (self.r_max_value - 1) * np.exp(-t_val)) - d_val = self.d_max_value / (1 + ((self.d_max_value / 1e-3) - 1) * np.exp(-(2 * t_val))) - t_val += float(self.t_delta) + # update r_max and d_max + t_val = K.get_value(self.t) + r_val = self.r_max_value / (1 + (self.r_max_value - 1) * np.exp(-t_val)) + d_val = self.d_max_value / (1 + ((self.d_max_value / 1e-3) - 1) * np.exp(-(2 * t_val))) + t_val += float(self.t_delta) - self.add_update([K.update(self.r_max, r_val), - K.update(self.d_max, d_val), - K.update(self.t, t_val)], x) + self.add_update([K.update(self.r_max, r_val), + K.update(self.d_max, d_val), + K.update(self.t, t_val)], inputs) - if self.mode == 0: - if sorted(reduction_axes) == range(K.ndim(x))[:-1]: + if training in {0, False}: + return x_normed + else: + def normalize_inference(): + if sorted(reduction_axes) == range(K.ndim(inputs))[:-1]: x_normed_running = K.batch_normalization( - x, self.running_mean, self.running_std, + inputs, self.running_mean, self.running_variance, self.beta, self.gamma, epsilon=self.epsilon) + + return x_normed_running else: # need broadcasting broadcast_running_mean = K.reshape(self.running_mean, broadcast_shape) - broadcast_running_std = K.reshape(self.running_std, broadcast_shape) + broadcast_running_std = K.reshape(self.running_variance, broadcast_shape) broadcast_beta = K.reshape(self.beta, broadcast_shape) broadcast_gamma = K.reshape(self.gamma, broadcast_shape) x_normed_running = K.batch_normalization( - x, broadcast_running_mean, broadcast_running_std, + inputs, broadcast_running_mean, broadcast_running_std, broadcast_beta, broadcast_gamma, epsilon=self.epsilon) - # pick the normalized form of x corresponding to the training phase - # for batch renormalization, inference time remains same as batchnorm - x_normed = K.in_train_phase(x_normed, x_normed_running) + return x_normed_running - elif self.mode == 1: - # sample-wise normalization - m = K.mean(x, axis=self.axis, keepdims=True) - std = K.sqrt(K.var(x, axis=self.axis, keepdims=True) + self.epsilon) - x_normed_batch = (x - m) / (std + self.epsilon) + # pick the normalized form of inputs corresponding to the training phase + # for batch renormalization, inference time remains same as batchnorm + x_normed = K.in_train_phase(x_normed, normalize_inference, training=training) - r_max_value = K.get_value(self.r_max) - r = std / (self.running_std + self.epsilon) - r = K.stop_gradient(K.clip(r, 1 / r_max_value, r_max_value)) - - d_max_value = K.get_value(self.d_max) - d = (m - self.running_mean) / (self.running_std + self.epsilon) - d = K.stop_gradient(K.clip(d, -d_max_value, d_max_value)) - - x_normed = ((x_normed_batch * r) + d) * self.gamma + self.beta - - # update r_max and d_max - t_val = K.get_value(self.t) - r_val = self.r_max_value / (1 + (self.r_max_value - 1) * np.exp(-t_val)) - d_val = self.d_max_value / (1 + ((self.d_max_value / 1e-3) - 1) * np.exp(-(2 * t_val))) - t_val += float(self.t_delta) - - self.add_update([K.update(self.r_max, r_val), - K.update(self.d_max, d_val), - K.update(self.t, t_val)], x) - - return x_normed + return x_normed def get_config(self): config = {'epsilon': self.epsilon, - 'mode': self.mode, 'axis': self.axis, - 'gamma_regularizer': regularizers.serialize(self.gamma_regularizer), - 'beta_regularizer': regularizers.serialize(self.beta_regularizer), + 'gamma_regularizer': initializers.serialize(self.gamma_regularizer), + 'beta_regularizer': initializers.serialize(self.beta_regularizer), + 'moving_mean_initializer': initializers.serialize(self.moving_mean_initializer), + 'moving_variance_initializer': initializers.serialize(self.moving_variance_initializer), + 'beta_constraint': constraints.serialize(self.beta_constraint), + 'gamma_constraint': constraints.serialize(self.gamma_constraint), 'momentum': self.momentum, 'r_max_value': self.r_max_value, 'd_max_value': self.d_max_value, diff --git a/tests/keras_contrib/layers/test_normalization.py b/tests/keras_contrib/layers/test_normalization.py index 29a32c7..ecd569a 100644 --- a/tests/keras_contrib/layers/test_normalization.py +++ b/tests/keras_contrib/layers/test_normalization.py @@ -18,21 +18,21 @@ input_shapes = [np.ones((10, 10)), np.ones((10, 10, 10))] @keras_test def basic_batchrenorm_test(): from keras import regularizers + layer_test(normalization.BatchRenormalization, - kwargs={'mode': 1, - 'gamma_regularizer': regularizers.l2(0.01), - 'beta_regularizer': regularizers.l2(0.01)}, input_shape=(3, 4, 2)) + layer_test(normalization.BatchRenormalization, - kwargs={'mode': 0}, + kwargs={'gamma_regularizer': regularizers.l2(0.01), + 'beta_regularizer': regularizers.l2(0.01)}, input_shape=(3, 4, 2)) @keras_test def test_batchrenorm_mode_0_or_2(): - for mode in [0, 2]: + for training in [1, 0]: model = Sequential() - norm_m0 = normalization.BatchRenormalization(mode=mode, input_shape=(10,), momentum=0.8) + norm_m0 = normalization.BatchRenormalization(input_shape=(10,), momentum=0.8) model.add(norm_m0) model.compile(loss='mse', optimizer='sgd') @@ -52,8 +52,8 @@ def test_batchrenorm_mode_0_or_2_twice(): # This is a regression test for issue #4881 with the old # batch normalization functions in the Theano backend. model = Sequential() - model.add(normalization.BatchRenormalization(mode=0, input_shape=(10, 5, 5), axis=1)) - model.add(normalization.BatchRenormalization(mode=0, input_shape=(10, 5, 5), axis=1)) + model.add(normalization.BatchRenormalization(input_shape=(10, 5, 5), axis=1)) + model.add(normalization.BatchRenormalization(input_shape=(10, 5, 5), axis=1)) model.compile(loss='mse', optimizer='sgd') X = np.random.normal(loc=5.0, scale=10.0, size=(20, 10, 5, 5)) @@ -64,7 +64,7 @@ def test_batchrenorm_mode_0_or_2_twice(): @keras_test def test_batchrenorm_mode_0_convnet(): model = Sequential() - norm_m0 = normalization.BatchRenormalization(mode=0, axis=1, input_shape=(3, 4, 4), momentum=0.8) + norm_m0 = normalization.BatchRenormalization(axis=1, input_shape=(3, 4, 4), momentum=0.8) model.add(norm_m0) model.compile(loss='mse', optimizer='sgd') @@ -79,27 +79,13 @@ def test_batchrenorm_mode_0_convnet(): assert_allclose(np.std(out, axis=(0, 2, 3)), 1.0, atol=1e-1) -@keras_test -def test_batchrenorm_mode_1(): - norm_m1 = normalization.BatchRenormalization(input_shape=(10,), mode=1) - norm_m1.build(input_shape=(None, 10)) - - for inp in [input_1, input_2, input_3]: - out = (norm_m1.call(K.variable(inp)) - norm_m1.beta) / norm_m1.gamma - assert_allclose(K.eval(K.mean(out)), 0.0, atol=1e-1) - if inp.std() > 0.: - assert_allclose(K.eval(K.std(out)), 1.0, atol=1e-1) - else: - assert_allclose(K.eval(K.std(out)), 0.0, atol=1e-1) - - @keras_test def test_shared_batchrenorm(): '''Test that a BN layer can be shared across different data streams. ''' # Test single layer reuse - bn = normalization.BatchRenormalization(input_shape=(10,), mode=0) + bn = normalization.BatchRenormalization(input_shape=(10,)) x1 = Input(shape=(10,)) bn(x1)