Why BatchNormalization fails in Keras

Question

I try to test ResNet approach on cifar10 dataset with the following python code:

# load data
(X_train, y_train), (X_test, y_test) = cifar10.load_data()
X_train = X_train.astype('float32')/256
X_test = X_test.astype('float32')/256
y_train = keras.utils.to_categorical(y_train)
y_test = keras.utils.to_categorical(y_test)

# build a model
def res_unit(x):
    x_shortcut = x

    x = Conv2D(16, (1, 1), padding='same')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)

    x = Conv2D(16, (3, 3), padding='same')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)

    x = Conv2D(32, (1, 1), padding='same')(x)
    x = BatchNormalization()(x)
    x = Add()([x_shortcut, x])

    x = Activation('relu')(x)

    return x


X_input = Input((X_train.shape[1:]))

X = Conv2D(32, (3, 3), padding='same')(X_input)
X = BatchNormalization()(X)
X = Activation('relu')(X)

X = res_unit(X)
X = res_unit(X)
X = res_unit(X)

X = Flatten()(X)
X = Dense(32)(X)
X = Activation('relu')(X)
X = Dense(10)(X)

X = Activation('softmax')(X)

model = Model(inputs=X_input, outputs=X)
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

# run model
model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=5)

which outputs the worst result for 10-class classification: 0.1 accuracy on train and validation.

But when I comment all lines with BatchNormalization, everything starts to be OK: I got 0.57 accuracy after first epoch and then it even rises up.

What is the problem with BatchNormalization in my case? Am I using it in a correct way?

Answer 1

I have played around with it and managed to get a working configuration. Necessary changes:

• axis=1 for BatchNormalization. Normally, this should be used when data_format="channels_first", but it worked for me despite TensorFlow backend

• Standardize the input:

  X_train = X_train * 2.0 - 1.0; X_test = X_test * 2.0 - 1.0

• Avoid batchnorm after the last convolution in res_unit.

The working solution looks like this:

# build a model                                                                                                                                                                                                                                                       
def res_unit(x, prefix=''):                                                                                                                                                                                                                                           
    x_shortcut = x                                                                                                                                                                                                                                                    

    x = Conv2D(16, (1, 1), padding='same', name='{}/{}'.format(prefix, 'conv1'))(x)                                                                                                                                                                                   
    x = BatchNormalization(name='{}/{}/bn'.format(prefix, 'conv1'), axis=1)(x)                                                                                                                                                                                        
    x = Activation(activation='relu', name='{}/{}/relu'.format(prefix, 'conv1'))(x)                                                                                                                                                                                   

    x = Conv2D(16, (3, 3), padding='same', name='{}/{}'.format(prefix, 'conv2'))(x)                                                                                                                                                                                   
    x = BatchNormalization(name='{}/{}/bn'.format(prefix, 'conv2'), axis=1)(x)                                                                                                                                                                                        
    x = Activation(activation='relu', name='{}/{}/relu'.format(prefix, 'conv2'))(x)                                                                                                                                                                                   

    x = Conv2D(32, (1, 1), padding='same', name='{}/{}'.format(prefix, 'conv3'))(x)                                                                                                                                                                                   
    x = Add(name='{}/add'.format(prefix))([x_shortcut, x])                                                                                                                                                                                                            
    x = Activation('relu', name='{}/relu'.format(prefix))(x)                                                                                                                                                                                                          

    return x                                                                                                                                                                                                                                                          


def main():                                                                                                                                                                                                                                                           
    (X_train, y_train), (X_test, y_test) = load_data()                                                                                                                                                                                                                
    X_train = X_train.astype('float32') / 255                                                                                                                                                                                                                         
    X_test = X_test.astype('float32') / 255                                                                                                                                                                                                                           
    X_train = X_train * 2.0 - 1.0                                                                                                                                                                                                                                     
    X_test = X_test * 2.0 - 1.0                                                                                                                                                                                                                                       
    y_train = keras_utils.to_categorical(y_train)                                                                                                                                                                                                                     
    y_test = keras_utils.to_categorical(y_test)                                                                                                                                                                                                                       

    print(X_train.min())                                                                                                                                                                                                                                              
    print(X_train.max())                                                                                                                                                                                                                                              

    X_input = Input((X_train.shape[1:]))                                                                                                                                                                                                                              

    X = Conv2D(32, (3, 3), padding='same', name='conv0')(X_input)                                                                                                                                                                                                     
    X = BatchNormalization(name='conv0/bn', axis=1)(X)                                                                                                                                                                                                                
    X = Activation('relu', name='conv0/relu')(X)                                                                                                                                                                                                                      

    X = res_unit(X, prefix='block1')                                                                                                                                                                                                                                  
    X = res_unit(X, prefix='block2')                                                                                                                                                                                                                                  
    X = res_unit(X, prefix='block3')                                                                                                                                                                                                                                  

    X = Flatten()(X)                                                                                                                                                                                                                                                  
    X = Dense(32)(X)                                                                                                                                                                                                                                                  
    X = Activation('relu')(X)                                                                                                                                                                                                                                         
    X = Dense(10)(X)                                                                                                                                                                                                                                                  

    X = Activation('softmax', name='softmax')(X)                                                                                                                                                                                                                      

    model = Model(inputs=X_input, outputs=X)                                                                                                                                                                                                                          
    model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

I've got 59.07% validation accuracy with it. However, sometimes training does not converge, it feels like it is very sensitive to the random init. Still don't know why making batchnorm work here is so hard...

Answer 2

I thank Serge who came up with a problem example (an issue) of using ResNet which many keras-API users believe that this is related to the BatchNormalization problem. I especially thank Dmytro Prylipko who suggested to modify the parameter "axis" when writing the line: x = BatchNormalization (axis = 3)(x) to x = BatchNormalization (axis = 1)(x) . Today I just tried your sugestion to my (similar) case in ResNet, and it seems to work. Previously, when I run a ResNet program, the training accuracy & loss are Ok (excellent), but the validation accuracy & loss behaved very ERRATIC. After I changed the paramater in the BatchNormalization --> from (axis = 3) to (axis = 1), the first 3 - 4 epochs still behaved erratically, but starting from the 5th Epoch, It "suddenly" became "stabilized". The accuracy & loss of the validation parts +/- "follow" those of the training parts. Again, thankyou.

By the way, previously, following another suggestion, at the beginning of the program I added 2 lines :

[1] from keras import backend as K ... and ... [2] K.set_learning_phase(1)

(the context of the above suggestion is because since the newer keras (2.1?), the user can no longer select the "mode" parameter in the BatchNormalizaton (either mode = 1 or mode = 2, to let the whole program knows whether it is a LEARNING-phase or a TEST/INFERENCE phase).

And actually few hourse ago, before reading this discussion, I was also "tempted" to try to make the same parameter change : (axis = 3) to (axis = 1), but then I changed my mind before really try that, because it doesn't seem to be making sense :-) .