Update BatchRenorm to Keras 2 API (#69)

* Update BatchRenorm to Keras 2 API

* update tests

* Remove mode 1 test (mode 1 doesnt exist anymore)
This commit is contained in:
Somshubra Majumdar
2017-04-15 09:43:55 -07:00
committed by Michael Oliver
parent 88c60d2459
commit 8ef69698b8
2 changed files with 136 additions and 149 deletions
+126 -125
View File
@@ -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,
@@ -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)