Use default (non-shared) quantization params for HardTanh

Author: skywall

backends/nxp/quantizer/neutron_quantizer.py

@@ -169,7 +169,7 @@ def partition_types(self):
         return [torch.ops.aten.relu_.default]
 
 
-class HardTanhPattern(SharedSpecPattern):
+class HardTanhPattern(QuantizationPattern):
     """
     Quantizer for HardTanh operator. Shared quantization spec is selected, as activation functions usually follows
     computation layer.
@@ -178,8 +178,23 @@ class HardTanhPattern(SharedSpecPattern):
     def partition_types(self):
         return [torch.ops.aten.hardtanh.default]
 
+    def get_anchors(
+            self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule]
+    ) -> PartitionAnchors | None:
+        node = fused_partition[0].nodes[-1]
+
+        return PartitionAnchors(
+            inputs=[(node, 0)],
+            weights=[],
+            biases=[],
+            output=[(node,)],
+        )
+
+    def replacement_op(self):
+        assert False
+
 
-class HardTanhInPlacePattern(SharedSpecPattern):
+class HardTanhInPlacePattern(QuantizationPattern):
     """
     Quantizer for HardTanh operator with param inplace=True. Shared quantization spec is selected, as activation
     functions usually follows computation layer.
@@ -188,6 +203,21 @@ class HardTanhInPlacePattern(SharedSpecPattern):
     def partition_types(self):
         return [torch.ops.aten.hardtanh_.default]
 
+    def get_anchors(
+            self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule]
+    ) -> PartitionAnchors | None:
+        node = fused_partition[0].nodes[-1]
+
+        return PartitionAnchors(
+            inputs=[(node, 0)],
+            weights=[],
+            biases=[],
+            output=[(node,)],
+        )
+
+    def replacement_op(self):
+        assert False
+
 
 class ReshapePattern(SharedSpecPattern):
     """

backends/nxp/tests/executorch_pipeline.py

@@ -28,6 +28,12 @@ def _quantize_model(model, calibration_inputs: list[tuple[torch.Tensor]]):
     return m
 
 
+def get_random_float_data(input_shapes: tuple[int] | list[tuple[int]]):
+    #TODO(Lukas): Replace with something more robust.
+    return (torch.randn(input_shapes),) if type(input_shapes) is tuple \
+        else tuple(torch.randn(input_shape) for input_shape in input_shapes)
+
+
 def to_quantized_edge_program(model: torch.nn.Module, input_shapes: tuple[int] | list[tuple[int]],
                               operators_not_to_delegate: list[str] = None, target="imxrt700",
                               neutron_converter_flavor="wrapper", remove_quant_io_ops=False)\
@@ -36,9 +42,7 @@ def to_quantized_edge_program(model: torch.nn.Module, input_shapes: tuple[int] |
         assert all([isinstance(input_shape, tuple) for input_shape in input_shapes]), ("For multiple inputs, provide"
                                                                                        " list[tuple[int]].")
 
-    random_tensors  = (torch.randn(input_shapes),) if type(input_shapes) is tuple \
-        else tuple(torch.randn(input_shape) for input_shape in input_shapes)
-    calibration_inputs = [random_tensors, random_tensors]
+    calibration_inputs = [get_random_float_data(input_shapes) for _ in range(4)]
     example_input = (torch.ones(input_shapes),) if type(input_shapes) is tuple \
         else tuple(torch.ones(input_shape) for input_shape in input_shapes)
 

backends/nxp/tests/ir/converter/node_converter/test_hardtanh_converter.py

@@ -22,15 +22,15 @@ class Relu6ConvBlock(torch.nn.Module):
     def __init__(self, conv_in_channels: int = 3, inplace: bool = False):
         super().__init__()
         self.block = torch.nn.Sequential(
-          torch.nn.Conv2d(in_channels=conv_in_channels, out_channels=64, kernel_size=(4, 4)),
-          torch.nn.ReLU6(inplace=inplace)
+            torch.nn.Conv2d(in_channels=conv_in_channels, out_channels=64, kernel_size=(4, 4)),
+            torch.nn.ReLU6(inplace=inplace)
         )
 
     def forward(self, x):
         return self.block(x)
 
 
-class CustomHardTanhBlock(torch.nn.Module):
+class ConvHardTanhBlock(torch.nn.Module):
     def __init__(self,
                  conv_in_channels: int = 3,
                  min_act_val: float = -1.,
@@ -70,12 +70,14 @@ def test_relu6_quant(mocker, input_shape: tuple[int], inplace: bool):
                         atol=1.)
 
 
-@pytest.mark.parametrize('input_shape', [(1, 3, 128, 128), (1, 3, 256, 256)])
+@pytest.mark.parametrize('input_shape', [(1, 3, 16, 16), (1, 3, 32, 32)])
 @pytest.mark.parametrize('activation_range', list(HardTanhConverter.supported_modes_map.keys()))
 @pytest.mark.parametrize('inplace', [True, False])
 def test_custom_hardtanh_quant(mocker, input_shape: tuple[int], activation_range: tuple[int, int], inplace: bool):
+    # TODO(Lukas): This test suffers from non-ideal testing random quantization, because we always use range <0,1>.
+    #  We should update this (decrease atol) when we have custom calibration dataset definition in place.
     min_val, max_val = activation_range
-    model = CustomHardTanhBlock(
+    model = ConvHardTanhBlock(
         conv_in_channels=input_shape[1],
         min_act_val=min_val,
         max_act_val=max_val,
@@ -98,4 +100,4 @@ def test_custom_hardtanh_quant(mocker, input_shape: tuple[int], activation_range
                         tflite_input_preprocess=ToChannelLastPreprocess(),
                         tflite_output_preprocess=ToChannelFirstPreprocess(),
                         input_data=input_data,
-                        atol=1.)
+                        atol=2.)