Merge pull request #1161 from JanFSchulte/transposefix

Bug fixes for channel-last conversions in pytorch
fastmachinelearning · Jan 7, 2025 · 3fa2902 · 3fa2902
2 parents 4b7e12d + 31219e3
commit 3fa2902
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/hls4ml/model/optimizer/passes/convert_to_channels_last.py b/hls4ml/model/optimizer/passes/convert_to_channels_last.py
@@ -97,12 +97,17 @@ def transform(self, model, node):
             if (
                 isinstance(node, Reshape)
                 and len(node.attributes['target_shape']) == 1
-                and not model.config.config['HLSConfig']['Model']['ChannelsLastConversion'] == "internal"
+                and not model.config.config['HLSConfig']['Model']['ChannelsLastConversion'] == "off"
             ):
                 previous_node = node.get_input_node(node.inputs[0])
                 input = previous_node.name
                 outshape = previous_node.get_output_variable().shape
 
+                if (model.config.config['IOType'] == 'io_stream') and len(outshape) == 3:
+                    raise Exception(
+                        'No 3D transpose available in io_stream, this model cannot be converted to channels-last'
+                    )
+
                 if len(outshape) == 2:
                     attributes = {'perm': [1, 0]}
                 else:

diff --git a/hls4ml/utils/config.py b/hls4ml/utils/config.py
@@ -283,7 +283,7 @@ def config_from_pytorch_model(
     default_precision='ap_fixed<16,6>',
     default_reuse_factor=1,
     channels_last_conversion='full',
-    transpose_outputs=True,
+    transpose_outputs=False,
     max_precision=None,
 ):
     """Create an HLS conversion config given the PyTorch model.

diff --git a/test/pytest/test_pytorch_api.py b/test/pytest/test_pytorch_api.py
@@ -498,7 +498,7 @@ def test_pooling(pooling, padds, backend):
     model.eval()
     pytorch_prediction = model(torch.Tensor(X_input)).detach().numpy()
 
-    config = config_from_pytorch_model(model, input_shape_forHLS)
+    config = config_from_pytorch_model(model, input_shape_forHLS, transpose_outputs=True)
     output_dir = str(test_root_path / f'hls4mlprj_pytorch_api_pooling_{pooling.__name__}_padds_{padds}_backend_{backend}')
     hls_model = convert_from_pytorch_model(model, hls_config=config, output_dir=output_dir, backend=backend)
     hls_model.compile()