update

Author: jmduarte

hls4ml/converters/keras/core.py

@@ -131,6 +131,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'] = not pytorch_layer.affine
 
     in_size = 1
     for dim in input_shapes[0][1:]:

hls4ml/model/layers.py

@@ -863,6 +863,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):
@@ -871,16 +873,11 @@ 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')
 
-        if gamma is None:
-            gamma = np.ones(mean.shape)
-        if beta is None:
-            beta = np.zeros(mean.shape)
-
         scale = gamma / np.sqrt(var + self.get_attr('epsilon'))
         bias = beta - scale * mean
 

test/pytest/test_batchnorm.py

@@ -23,14 +23,14 @@ def data():
 @pytest.fixture(scope='module')
 def model(request):
     model = Sequential()
-    model.add(BatchNormalization(input_shape=(in_shape,), center=request.param[0], scale=request.param[1]))
+    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, True], [True, False], [False, True], [False, False]), indirect=True)
+@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>'