conv

Author: xadupre

_unittests/ut_reference/test_torch_onnx_evaluator.py

@@ -1069,6 +1069,41 @@ def test_loop(self):
             model, torch.tensor(5, dtype=torch.int64), torch.tensor(1, dtype=torch.bool)
         )
 
+    def test_conv(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node(
+                        "Conv",
+                        ["X", "W", "B"],
+                        ["Y"],
+                        pads=[1, 1, 1, 1],
+                        dilations=[1, 1],
+                        strides=[2, 2],
+                    )
+                ],
+                "g",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, [None, None, None, None]),
+                    oh.make_tensor_value_info("W", TFLOAT, [None, None, None, None]),
+                    oh.make_tensor_value_info("B", TFLOAT, [None, None, None, None]),
+                ],
+                [oh.make_tensor_value_info("Y", TFLOAT, [None, None, None, None])],
+            ),
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        sH, sW = 5, 6
+        i = sH // 2
+        j = sW // 2
+        X = torch.zeros((1, 1, sH, sW), dtype=torch.float32)
+        X[0, 0, i, j] = 1.0
+        W = torch.zeros((1, 1, 3, 3), dtype=torch.float32)
+        W[0, 0, :, :] = torch.minimum(
+            2 ** torch.arange(9).reshape((3, -1)), torch.tensor([256])
+        )
+        B = torch.tensor([[[[0]]]], dtype=torch.float32)
+        self._finalize_test(model, X, W, B, use_ort=True)
+
 
 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 LayerNormalization_17, Softmax_13, Tanh_6
+from .nn_ops import Conv_11, LayerNormalization_17, Softmax_13, Tanh_6
 from .other_ops import Cast_6, CastLike_15, Concat_1, Transpose_1, Trilu_14, Where_9
 from .reduce_ops import ReduceMax_18, ReduceMean_18, ReduceMin_18, ReduceSum_18
 from .sequence_ops import ConcatFromSequence_11, SequenceEmpty_11, SequenceInsert_11

onnx_diagnostic/reference/torch_ops/_op_run.py

@@ -261,7 +261,21 @@ def get_attribute_ints(
         :return: value
         """
         att = self._find_attribute(node, name)
-        return default_value if att is None else tuple(att.ints)
+        return default_value if att is None else tuple(map(int, att.ints))
+
+    def get_attribute_string(
+        self, node: onnx.NodeProto, name: str, default_value: Optional[str] = None
+    ) -> Optional[Tuple[int, ...]]:
+        """
+        Returns an attribute as a tuple of ints.
+
+        :param node: NodeProto
+        :param name: name
+        :param default_value: default_value
+        :return: value
+        """
+        att = self._find_attribute(node, name)
+        return default_value if att is None else att.s.decode("utf-8")
 
     def get_attribute_tensor(self, node: onnx.NodeProto, name: str) -> Optional[torch.Tensor]:
         """

onnx_diagnostic/reference/torch_ops/nn_ops.py

@@ -5,6 +5,49 @@
 from . import OpRun, OpRunTensor
 
 
+class Conv_11(OpRun):
+    "Conv"
+
+    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.dilations = self.get_attribute_ints(node, "dilations", None)
+        self.group = self.get_attribute_int(node, "group", 1)
+        self.kernel_shape = self.get_attribute_int(node, "kernel_shape", [])
+        self.pads = self.get_attribute_ints(node, "pads", None)
+        self.strides = self.get_attribute_ints(node, "strides", None)
+
+    def run(self, x, w, b=None):
+        kernel_shape = self.kernel_shape or w.shape[2:]
+        assert (
+            tuple(kernel_shape) == w.shape[2:]
+        ), 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}"
+        if b is None:
+            bias = None
+        else:
+            bias = b.tensor.squeeze()
+            if not bias.shape:
+                bias = bias.unsqueeze(0)
+        return OpRunTensor(
+            torch.nn.functional.conv2d(
+                x.tensor,
+                w.tensor,
+                bias=bias,
+                stride=tuple(strides),
+                padding=pads[0],
+                dilation=tuple(dilations),
+                groups=self.group,
+            )
+        )
+
+
 class LayerNormalization_17(OpRun):
     "LayerNormalization"