add averagepool

Author: xadupre

_unittests/ut_reference/test_torch_onnx_evaluator.py

@@ -1185,6 +1185,63 @@ def test_scatternd_2d(self):
                     use_ort=True,
                 )
 
+    def test_averagepool_1d(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node(
+                        "AveragePool",
+                        inputs=["x"],
+                        outputs=["y"],
+                        kernel_shape=[2],
+                    )
+                ],
+                "ut",
+                [oh.make_tensor_value_info("x", TFLOAT, [None, None, None])],
+                [oh.make_tensor_value_info("y", TFLOAT, [None, None, None])],
+            ),
+            opset_imports=[oh.make_opsetid("", 18)],
+            ir_version=10,
+        )
+        self._finalize_test(model, torch.rand((1, 3, 32), dtype=torch.float32))
+
+    def test_averagepool_2d(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node(
+                        "AveragePool",
+                        inputs=["x"],
+                        outputs=["y"],
+                        kernel_shape=[5, 5],
+                        pads=[2, 2, 2, 2],
+                    )
+                ],
+                "ut",
+                [oh.make_tensor_value_info("x", TFLOAT, [None, None, None, None])],
+                [oh.make_tensor_value_info("y", TFLOAT, [None, None, None, None])],
+            ),
+            opset_imports=[oh.make_opsetid("", 18)],
+            ir_version=10,
+        )
+        self._finalize_test(
+            model,
+            torch.tensor(
+                [
+                    [
+                        [
+                            [1, 2, 3, 4, 5],
+                            [6, 7, 8, 9, 10],
+                            [11, 12, 13, 14, 15],
+                            [16, 17, 18, 19, 20],
+                            [21, 22, 23, 24, 25],
+                        ]
+                    ]
+                ],
+                dtype=torch.float32,
+            ),
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -17,7 +17,7 @@
 )
 from .controlflow_ops import If_1, Loop_16
 from .generator_ops import Range_11
-from .nn_ops import Conv_11, LayerNormalization_17, Softmax_13, Tanh_6
+from .nn_ops import AveragePool_11, Conv_11, LayerNormalization_17, Softmax_13, Tanh_6
 from .other_ops import (
     Cast_6,
     CastLike_15,

onnx_diagnostic/reference/torch_ops/nn_ops.py

@@ -5,6 +5,45 @@
 from . import OpRun, OpRunTensor
 
 
+class AveragePool_11(OpRun):
+    "AveragePool"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        self.auto_pad = self.get_attribute_string(node, "auto_pad", "NOTSET")
+        self.ceil_mode = bool(self.get_attribute_int(node, "ceil_mode", 0))
+        self.count_include_pad = bool(self.get_attribute_int(node, "count_include_pad", 0))
+        self.dilations = self.get_attribute_ints(node, "dilations", None)
+        self.kernel_shape: Tuple[int, ...] = (
+            self.get_attribute_ints(node, "kernel_shape") or tuple()
+        )
+        self.pads = self.get_attribute_ints(node, "pads", None)
+        self.strides = self.get_attribute_ints(node, "strides", None)
+
+    def run(self, x):
+        kernel_shape = self.kernel_shape
+        dilations = self.dilations or [1 for _ in x.shape[2:]]
+        strides = self.strides or [1 for _ in x.shape[2:]]
+        pads = self.pads or ([0 for _ in x.shape[2:]] * 2)
+        assert (
+            self.auto_pad == "NOTSET"
+        ), f"conv not implemented for auto_pad={self.auto_pad!r}"
+        assert len(set(pads)) == 1, f"conv not implemented for pads={pads}"
+        assert set(dilations) == {1}, f"conv not implemented for dilations={dilations}"
+        avg_pool = getattr(torch.nn.functional, f"avg_pool{len(kernel_shape)}d")
+        return OpRunTensor(
+            avg_pool(
+                x.tensor,
+                kernel_size=tuple(kernel_shape),
+                stride=tuple(strides),
+                padding=pads[0],
+                ceil_mode=self.ceil_mode,
+                count_include_pad=self.count_include_pad,
+                # dilation=tuple(dilations),
+            )
+        )
+
+
 class Conv_11(OpRun):
     "Conv"
 
@@ -22,15 +61,15 @@ def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
     def run(self, x, w, b=None):
         kernel_shape = self.kernel_shape or w.shape[2:]
         assert (
-            tuple(kernel_shape) == w.shape[2:]
+            tuple(kernel_shape) == w.shape[-len(kernel_shape) :]
         ), f"conv not implemented for kernel_shape={kernel_shape} and w.shape={w.shape}"
         dilations = self.dilations or [1 for _ in x.shape[2:]]
         strides = self.strides or [1 for _ in x.shape[2:]]
         pads = self.pads or ([0 for _ in x.shape[2:]] * 2)
         assert (
             self.auto_pad == "NOTSET"
         ), f"conv not implemented for auto_pad={self.auto_pad!r}"
-        assert len(set(pads)) == 1, f"conv not implemented for dilations={pads}"
+        assert len(set(pads)) == 1, f"conv not implemented for pads={pads}"
         if b is None:
             bias = None
         else: