ADD support for PQBatchNorm1d in pytorch

Author: Enrico Lupi

hls4ml/converters/pytorch/pquant.py

@@ -111,7 +111,8 @@ def handler(operation, layer_name, input_names, input_shapes, node, class_object
 
 
 parse_pqlinear_layer = make_pquant_handler(parse_linear_layer, 'PQDense', 'PQLinear')
-parse_pqbatchnorm_layer = make_pquant_handler(parse_batchnorm_layer, 'PQBatchNorm2d')
+parse_pqbatchnorm1d_layer = make_pquant_handler(parse_batchnorm_layer, 'PQBatchNorm1d')
+parse_pqbatchnorm2d_layer = make_pquant_handler(parse_batchnorm_layer, 'PQBatchNorm2d')
 parse_pqconv1d_layer = make_pquant_handler(parse_conv1d_layer, 'PQConv1d')
 parse_pqconv2d_layer = make_pquant_handler(parse_conv2d_layer, 'PQConv2d')
 parse_pqpool1d_layer = make_pquant_handler(parse_pooling_layer, 'PQAvgPool1d', 'AvgPool1d')

test/pytest/test_pquant_pytorch.py

@@ -3,17 +3,17 @@
 
 import numpy as np
 import pytest
-from pquant.activations import PQActivation
-from pquant.core.finetuning import TuningConfig
-from pquant.core.utils import get_default_config
-from pquant.layers import PQAvgPool1d, PQAvgPool2d, PQBatchNorm2d, PQConv1d, PQConv2d, PQDense
 
 from hls4ml.converters import convert_from_pytorch_model
 from hls4ml.utils import config_from_pytorch_model
 
 os.environ['KERAS_BACKEND'] = 'torch'
 import torch  # noqa: E402
 import torch.nn as nn  # noqa: E402
+from pquant.activations import PQActivation  # noqa: E402
+from pquant.core.finetuning import TuningConfig  # noqa: E402
+from pquant.core.utils import get_default_config  # noqa: E402
+from pquant.layers import PQAvgPool1d, PQAvgPool2d, PQBatchNorm1d, PQBatchNorm2d, PQConv1d, PQConv2d, PQDense  # noqa: E402
 
 test_path = Path(__file__).parent
 
@@ -125,9 +125,9 @@ def get_shape(model: nn.Module, batch_size: int = 1, default_length: int = 32, d
         case PQAvgPool2d():
             # (N, C, H, W)
             return (batch_size, 1, *default_hw)
-        # case PQBatchNorm1d():
-        #    # (N, num_features, L)
-        #    return (batch_size, layer.num_features, *default_length)
+        case PQBatchNorm1d():
+            # (N, num_features, L)
+            return (batch_size, layer.num_features, default_length)
         case PQBatchNorm2d():
             # (N, num_features, H, W)
             return (batch_size, layer.num_features, *default_hw)
@@ -159,6 +159,7 @@ def get_shape(model: nn.Module, batch_size: int = 1, default_length: int = 32, d
         'PQConv2d(2, 3, kernel_size=(3,3), padding=0, bias=False)',
         'PQConv2d(2, 3, kernel_size=(3,3), padding=0, stride=2)',
         'PQConv2d(2, 3, kernel_size=(3,3), padding=1, stride=2)',
+        'PQBatchNorm1d(3)',
         'PQBatchNorm2d(3)',
         'PQAvgPool1d(2, padding=1)',
         'PQAvgPool1d(2, padding=0)',