Keras: Batch normalization(axis=1) has rank 0

Question

I am trying to create a neural network with input (batch, channels(3), 64, 32) In Keras but I have errors using Batch normalization (Since the error starts from the mentioned layer I opted to isolate the part that causes the the error) The model starts like this:

input_imgC = Input(shape=(X.shape[1], X.shape[2], X.shape[3]))
x = Conv2D(32, (5, 5), activation='relu', padding='same', data_format="channels_first")(input_imgC) 
out = BatchNormalization(axis=1)(x)

I receive the following exception:

Shape must be rank 1 but is rank 0 for 'batch_normalization_1/cond/Reshape_4' (op: 'Reshape') with input shapes: [1,32,1,1], [].

I thought the data was badly converted from the convolutional layer, and that's why I tried the model without batch normalization. The model:

input_imgC = Input(shape=(X.shape[1], X.shape[2], X.shape[3]))
out = Conv2D(32, (5, 5), activation='relu', padding='same', data_format="channels_first")(input_imgC)

Produced correctly the following summary:

Layer (type) Output Shape Param #
================================================================= > input_1 (InputLayer) (None, 3, 64, 32) 0

conv2d_1 (Conv2D) (None, 32, 64, 32) 2432
================================================================= Total params: 2,432 Trainable params: 2,432 Non-trainable params: 0

I know I can use

out = BatchNormalization(axis=-1)(x)

with the model input as (batch, 64, 32, channels(3)) and it will work (I already tried it) but I need this configuration of channels at the beginning in order to test the model with a package that shows the saliency of an object.

Does any one have an Idea of what to do here?

Daniel Möller Daniel Möller · Accepted Answer · 2018-12-05T15:29:11

I don't know what the problem is, it seems a Keras bug (updating the version or searching for issues on GitHub might show hints about this).

My suggestion is that you permute dimensions at the beginning and use everything as channels last, as usual.

At the end, if your outputs are images, restore the dimensions to the order you want.

input_imgC = Input(shape=(X.shape[1], X.shape[2], X.shape[3]))
x = Lambda(lambda inputs: K.permute_dimensions(inputs,(0,2,3,1)))(input_imgC)

x = Conv2D(32, (5, 5), activation='relu', padding='same')(x) 
x = BatchNormalization()(x)

out = .......
out = Lambda(lambda inputs: K.permute_dimensions(inputs,(0,3,1,2)))(out)

Hint:

Ideally, I'd use batch normalization before relu. This turns relu into a safe activation that will very rarely get frozen, which is not true when used without care.

x = Conv2D(..., activation='linear')(x) #or no activation
x = BatchNormalization()(x)
x = Activation('relu')(x)

There are some advantages of doing this with other activations such as sigmoid and tanh, to, as the batch normalization may help escaping the saturated zones.

Keras: Batch normalization(axis=1) has rank 0

1 Answers