Merge pull request #754 from fastmachinelearning/batchnorm_fix

Fix for BatchNormalization layers with `center=False` or `scale=False`

Author: jmitrevs

hls4ml/converters/keras/core.py

@@ -98,6 +98,9 @@ def parse_batchnorm_layer(keras_layer, input_names, input_shapes, data_reader):
     elif len(input_shapes[0]) == 4:
         layer['n_filt'] = input_shapes[0][3]
 
+    layer['use_gamma'] = keras_layer['config']['scale']
+    layer['use_beta'] = keras_layer['config']['center']
+
     return layer, [shape for shape in input_shapes[0]]
 
 

hls4ml/converters/pytorch/core.py

@@ -65,6 +65,7 @@ def parse_batchnorm_layer(pytorch_layer, layer_name, input_shapes, data_reader,
 
     # batchnorm para
     layer['epsilon'] = pytorch_layer.eps
+    layer['use_gamma'] = layer['use_beta'] = pytorch_layer.affine
 
     in_size = 1
     for dim in input_shapes[0][1:]:

hls4ml/model/layers.py

@@ -541,7 +541,9 @@ def _get_folded_weights(self):
 
         # wrap conv kernel and bias with bn parameters
         folded_kernel = inv * kernel
-        folded_bias = inv * (bias - moving_mean) + beta
+        folded_bias = inv * (bias - moving_mean)
+        if beta is not None:
+            folded_bias += beta
 
         return [folded_kernel, folded_bias]
 
@@ -832,6 +834,8 @@ class BatchNormalization(Layer):
         WeightAttribute('bias'),
         TypeAttribute('scale'),
         TypeAttribute('bias'),
+        Attribute('use_gamma', value_type=bool, default=True),
+        Attribute('use_beta', value_type=bool, default=True),
     ]
 
     def initialize(self):
@@ -840,13 +844,13 @@ def initialize(self):
         dims = inp.dim_names
         self.add_output_variable(shape, dims)
 
-        gamma = self.model.get_weights_data(self.name, 'gamma')
-        beta = self.model.get_weights_data(self.name, 'beta')
+        gamma = self.model.get_weights_data(self.name, 'gamma') if self.get_attr('use_gamma') else 1
+        beta = self.model.get_weights_data(self.name, 'beta') if self.get_attr('use_beta') else 0
         mean = self.model.get_weights_data(self.name, 'moving_mean')
         var = self.model.get_weights_data(self.name, 'moving_variance')
 
         scale = gamma / np.sqrt(var + self.get_attr('epsilon'))
-        bias = beta - gamma * mean / np.sqrt(var + self.get_attr('epsilon'))
+        bias = beta - scale * mean
 
         self.add_weights_variable(name='scale', var_name='s{index}', data=scale)
         self.add_weights_variable(name='bias', var_name='b{index}', data=bias)

test/pytest/test_batchnorm.py

@@ -21,21 +21,24 @@ def data():
 
 
 @pytest.fixture(scope='module')
-def model():
+def model(request):
     model = Sequential()
-    model.add(BatchNormalization(input_shape=(in_shape,)))
+    model.add(BatchNormalization(input_shape=(in_shape,), center=request.param, scale=request.param))
     model.compile()
     return model
 
 
 @pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
 @pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus'])
+@pytest.mark.parametrize('model', [True, False], indirect=True)
 def test_batchnorm(model, data, backend, io_type):
 
     default_precision = 'ac_fixed<32, 1, true>' if backend == 'Quartus' else 'ac_fixed<32, 1>'
 
+    center = model.layers[0].center
+    scale = model.layers[0].scale
     config = hls4ml.utils.config_from_keras_model(model, default_precision=default_precision, granularity='name')
-    output_dir = str(test_root_path / f'hls4mlprj_batchnorm_{backend}_{io_type}')
+    output_dir = str(test_root_path / f'hls4mlprj_batchnorm_{backend}_{io_type}_center{center}_scale{scale}')
     hls_model = hls4ml.converters.convert_from_keras_model(
         model, backend=backend, hls_config=config, io_type=io_type, output_dir=output_dir
     )