Bug Fix for Operand Shape Mismatch in BatchNorm Fusion (PyTorch) (#1045)

* fixed operand dimension mismatch error in bn_fuse.py

* moved test file to test/

* added channels_last_conversion to config

* updating PR to match contribution guidelines. pre-commit has been run and test case has been moved from standalone file to existing pytests

* fix shape of bias tensor in pytorch if zero, add additional batchnorm tests

* consistent uses of '

* reverting changes in test_batchnorm_pytorch.py from merge

* reverting changes in test_batchnorm_pytorch.py from merge

---------

Co-authored-by: Jan-Frederik Schulte <jschulte@cern.ch>

Author: sei-rquartiano

hls4ml/model/layers.py

@@ -258,7 +258,13 @@ def add_bias(self, quantizer=None):
         precision = None
         type_name = None
         if data is None:
-            data = np.zeros(self.get_output_variable().shape[-1])
+            if 'data_format' in self.attributes:
+                if self.attributes['data_format'] == 'channels_first':
+                    data = np.zeros(self.get_output_variable().shape[0])
+                elif self.attributes['data_format'] == 'channels_last':
+                    data = np.zeros(self.get_output_variable().shape[-1])
+            else:
+                data = np.zeros(self.get_output_variable().shape[-1])
             precision = IntegerPrecisionType(width=1, signed=False)
             type_name = 'bias{index}_t'
             quantizer = None  # Don't quantize non-existant bias

test/pytest/test_batchnorm_pytorch.py

@@ -13,13 +13,22 @@
 atol = 5e-3
 
 
-@pytest.fixture(scope='module')
+@pytest.fixture
 def data():
     np.random.seed(0)
     X = np.random.rand(100, in_shape)
     return X
 
 
+@pytest.fixture(scope='module')
+def fusion_data():
+    n_batch = 2
+    n_in = 2
+    size_in_height = 32
+    X = np.random.rand(n_batch, n_in, size_in_height)
+    return X
+
+
 @pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
 @pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus', 'Catapult'])
 def test_batchnorm(data, backend, io_type):
@@ -41,3 +50,82 @@ def test_batchnorm(data, backend, io_type):
     pytorch_prediction = model(torch.Tensor(data)).detach().numpy()
     hls_prediction = hls_model.predict(data)
     np.testing.assert_allclose(pytorch_prediction, hls_prediction, rtol=0, atol=atol, verbose=True)
+
+
+atol = 5e-2
+
+
+class BatchNorm_w_Fusion(nn.Module):
+    def __init__(self, filters, momentum):
+        super().__init__()
+        self.conv1 = nn.Conv1d(
+            int(filters),
+            filters,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias=False,
+        )
+        self.bn1 = nn.BatchNorm1d(filters)
+        self.relu1 = nn.ReLU()
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu1(x)
+        return x
+
+
+@pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus', 'Catapult'])
+def test_batchnorm_fusion(fusion_data, backend, io_type):
+    n_in = 2
+    momentum = 0.99
+    size_in_height = 32
+    filters = n_in
+
+    # see above for model definition
+    model = BatchNorm_w_Fusion(filters, momentum)
+    # Important to set model to eval to fix batchnorm behavior
+    model.eval()
+    # generating config
+    pytorch_prediction = model(torch.Tensor(fusion_data)).detach().numpy()
+
+    # We do not have an implementation of a transpose for io_stream, need to transpose inputs and outputs outside of hls4ml
+    if io_type == 'io_stream':
+        fusion_data = np.ascontiguousarray(fusion_data.transpose(0, 2, 1))
+        config = hls4ml.utils.config_from_pytorch_model(model, channels_last_conversion='internal', transpose_outputs=False)
+    else:
+        config = hls4ml.utils.config_from_pytorch_model(model, channels_last_conversion='full', transpose_outputs=True)
+
+    config['Model']['Strategy'] = 'Resource'
+
+    default_precision = 'ac_fixed<32, 1, true>' if backend == 'Quartus' else 'ac_fixed<32, 1>'
+
+    config['Model']['Precision'] = default_precision
+
+    # conversion
+    output_dir = str(test_root_path / f'hls4mlprj_block_{backend}_{io_type}')
+    hls_model = hls4ml.converters.convert_from_pytorch_model(
+        model,
+        (None, n_in, size_in_height),
+        hls_config=config,
+        output_dir=output_dir,
+        backend=backend,
+        io_type=io_type,
+    )
+
+    # compiling model
+    hls_model.compile()
+
+    if io_type == 'io_stream':
+        hls_prediction = np.transpose(
+            np.reshape(
+                hls_model.predict(fusion_data),
+                (pytorch_prediction.shape[0], pytorch_prediction.shape[2], pytorch_prediction.shape[1]),
+            ),
+            (0, 2, 1),
+        )
+    else:
+        hls_prediction = np.reshape(hls_model.predict(fusion_data), pytorch_prediction.shape)
+    np.testing.assert_allclose(pytorch_prediction, hls_prediction, rtol=0, atol=atol, verbose=True)