Remove transpose of input if n_chan=1

Author: vloncar

hls4ml/model/optimizer/__init__.py

@@ -35,8 +35,9 @@
     [
         'channels_last_converter',
         'remove_transpose_before_flatten',
+        'remove_nop_transpose',
+        'remove_single_channel_transpose',
         'fuse_bias_add',
-        'remove_useless_transpose',
         'expand_layer_group',
         'output_rounding_saturation_mode',
         'qkeras_factorize_alpha',

hls4ml/model/optimizer/passes/transpose_opt.py

@@ -1,21 +1,47 @@
-from hls4ml.model.layers import Transpose
+from hls4ml.model.layers import Input, Transpose
 from hls4ml.model.optimizer import OptimizerPass
 
 
-class RemoveUselessTranspose(OptimizerPass):
+class RemoveNopTranspose(OptimizerPass):
+    """
+    Remove a transpose layer if it doesn't do anything to a 1D array. i.e, 1D input and perm = [0]
+    """
+
     def match(self, node):
         is_match = isinstance(node, Transpose) and node.get_attr('perm') == [0]  # Useless transpose
         return is_match
 
     def transform(self, model, node):
-        """
-        Remove a transpose layer if it doesn't do anything. i.e 1D input and perm = [0]
-        """
-        print(f"Unnessary {node.name} in the model, optimizing ...")
+        print(f'Unnecessary transpose node ({node.name}) detected, optimizing ...')
         if not node.get_output_nodes():
-            print(f"WARNING: {node.name} is the output layer! No rewiring performed.")
+            print(f'WARNING: {node.name} is the output layer! No rewiring performed.')
             model.remove_node(node, rewire=False)  # Don't rewire if there is no output layer
         else:
             model.remove_node(node, rewire=True)
 
         return True
+
+
+class RemoveSingleChannelTranspose(OptimizerPass):
+    """
+    Remove transpose of inputs if the number of channels is 1 as for io_parallel this doesn't affect the array
+    representation used
+    """
+
+    def match(self, node):
+        if node.model.config.get_config_value('IOType') != 'io_parallel':
+            return False
+
+        return (
+            isinstance(node, Transpose)
+            and isinstance(node.get_input_node(), Input)
+            and node.get_input_variable().shape[0] == 1
+        )
+
+    def transform(self, model, node):
+        # Adjust the input shape and remove the Transpose node
+        input_var = node.get_input_variable()
+        input_var.shape.append(input_var.shape.pop(0))
+        model.remove_node(node)
+
+        return True

test/pytest/test_pytorch_api.py

@@ -740,3 +740,72 @@ def test_skipped_layers(backend, io_type):
     hls_prediction = hls_model.predict(hls_input).flatten()
 
     np.testing.assert_allclose(hls_prediction, pytorch_prediction, rtol=0, atol=5e-2)
+
+
+@pytest.mark.parametrize('backend', ['Vivado', 'Quartus'])
+@pytest.mark.parametrize('io_type', ['io_parallel'])  # Only io_parallel for now
+@pytest.mark.parametrize('tensor_rank', [2, 3])
+def test_remove_transpose(backend, io_type, tensor_rank):
+    class TestModel(nn.Module):
+        def __init__(self, tensor_rank):
+            super().__init__()
+            if tensor_rank == 2:
+                self.conv1 = nn.Conv1d(in_channels=1, out_channels=4, kernel_size=3, bias=False)
+                self.relu1 = nn.ReLU()
+                self.flatten = nn.Flatten()
+                self.fc1 = nn.Linear(in_features=4 * 6, out_features=5, bias=False)
+                self.relu2 = nn.ReLU()
+            else:
+                self.conv1 = nn.Conv2d(in_channels=1, out_channels=4, kernel_size=3, bias=False)
+                self.relu1 = nn.ReLU()
+                self.flatten = nn.Flatten()
+                self.fc1 = nn.Linear(in_features=4 * 6 * 6, out_features=5, bias=False)
+                self.relu2 = nn.ReLU()
+
+        def forward(self, x):
+            # In the hls4ml model, there should be a Transpose node on the input tensor before conv1
+            x = self.conv1(x)
+            x = self.relu1(x)
+            x = self.flatten(x)  # This should result in a Transpose node that we aim to remove
+            x = self.fc1(x)
+            x = self.relu2(x)
+            return x
+
+    model = TestModel(tensor_rank=tensor_rank)
+    if tensor_rank == 2:
+        input_shape = (1, 8)
+        input_tensor = torch.randn(10, 1, 8)
+        hls_input = np.ascontiguousarray(torch.permute(input_tensor, (0, 2, 1)).detach().numpy())
+    else:
+        input_shape = (1, 8, 8)
+        input_tensor = torch.randn(10, 1, 8, 8)
+        hls_input = np.ascontiguousarray(torch.permute(input_tensor, (0, 2, 3, 1)).detach().numpy())
+
+    batch_input_shape = (None,) + input_shape
+    config = config_from_pytorch_model(
+        model,
+        default_precision='ap_fixed<32,16>',
+        inputs_channel_last=False,  # Crucial for testing if the first Transpose was removed
+        transpose_outputs=False,
+    )
+    output_dir = str(test_root_path / f'hls4mlprj_pytorch_api_transpose_nop_{tensor_rank}d_{backend}_{io_type}')
+    hls_model = convert_from_pytorch_model(
+        model,
+        batch_input_shape,
+        hls_config=config,
+        output_dir=output_dir,
+        io_type=io_type,
+        backend=backend,
+    )
+
+    hls_model.compile()
+
+    # Test optimizers removed the two Transpose layers
+    transpose_layers = [layer for layer in list(hls_model.get_layers()) if layer.class_name == 'Transpose']
+    assert len(transpose_layers) == 0
+
+    # Test predictions match
+    pytorch_prediction = model(input_tensor).detach().numpy().flatten()
+    hls_prediction = hls_model.predict(hls_input).flatten()
+
+    np.testing.assert_allclose(hls_prediction, pytorch_prediction, rtol=0, atol=5e-2)