keras: Error with BatchNormalization while setting the default float type to 'float16'

Hi,

I am trying to train a basic network on Keras with a float16 precision. However it looks like there is a bug with BatchNormalization.

For example, if I run the simple following code:

import keras
keras.backend.set_floatx('float16')

input_layer = keras.layers.Input(shape=(16,16,3))
x = keras.layers.BatchNormalization(axis=3)(input_layer) # <<Fails here
x = keras.layers.Conv2D(32, (3,3))(x)

I have the following error (entire output error at the end of this message):

TypeError: Value passed to parameter ‘scale’ has DataType float16 not in list of allowed values: float32

If I suppress all BatchNormalization layers and only leave Conv2D, Activation, MaxPooling and Dense, I have no error and the training runs fine. Same if I switch ‘float16’ to ‘float32’

I use Keras 2.1.4 with Tensorflow 1.4.0 as a backend, on a Linux machine with python 3.5.2.

Thank you for your help.

Traceback (most recent call last): File “test_BN_fp16.py”, line 6, in <module> x = keras.layers.BatchNormalization(axis=3)(input_layer) # <<Fails here File “/usr/local/lib/python3.5/dist-packages/keras/engine/topology.py”, line 617, in call output = self.call(inputs, **kwargs) File “/usr/local/lib/python3.5/dist-packages/keras/layers/normalization.py”, line 181, in call epsilon=self.epsilon) File “/usr/local/lib/python3.5/dist-packages/keras/backend/tensorflow_backend.py”, line 1824, in normalize_batch_in_training epsilon=epsilon) File “/usr/local/lib/python3.5/dist-packages/keras/backend/tensorflow_backend.py”, line 1799, in _fused_normalize_batch_in_training data_format=tf_data_format) File “/usr/local/lib/python3.5/dist-packages/tensorflow/python/ops/nn_impl.py”, line 831, in fused_batch_norm name=name) File “/usr/local/lib/python3.5/dist-packages/tensorflow/python/ops/gen_nn_ops.py”, line 2296, in _fused_batch_norm_v2 is_training=is_training, name=name) File “/usr/local/lib/python3.5/dist-packages/tensorflow/python/framework/op_def_library.py”, line 609, in _apply_op_helper param_name=input_name) File “/usr/local/lib/python3.5/dist-packages/tensorflow/python/framework/op_def_library.py”, line 60, in _SatisfiesTypeConstraint ", ".join(dtypes.as_dtype(x).name for x in allowed_list))) TypeError: Value passed to parameter ‘scale’ has DataType float16 not in list of allowed values: float32