feature: added Split operator (#81)

* added Split operator

* added tests for Split

* added printvisitor for Split

dnnv/nn/converters/onnx.py

@@ -497,16 +497,11 @@ def visit_OutputSelect(self, operation: operations.OutputSelect) -> onnx.NodePro
         idx = self.op_counts[op_type] = self.op_counts[op_type] + 1
         opname = f"{op_type}_{idx}"
 
-        if operation.index != 0:
-            raise NotImplementedError(
-                "Support for operations with multiple ouputs is not yet implemented."
-            )
-
         op = self._to_onnx_proto(operation.operation, f"{opname}.operation")
 
         node = onnx.helper.make_node(
             "Identity",
-            inputs=[op.name],
+            inputs=[op.output[operation.index]],
             outputs=[opname],
             name=opname,
         )
@@ -559,6 +554,28 @@ def visit_Sigmoid(self, operation: operations.Sigmoid) -> onnx.NodeProto:
 
         return node
 
+    def visit_Split(self, operation: operations.Split) -> onnx.NodeProto:
+        op_type = str(operation)
+        # TODO: split attribute is optional. Edits to nn/parser/onnx.py required.
+        assert operation.split is not None
+        idx = self.op_counts["Split"] = self.op_counts["Split"] + 1
+        opname = f"Split_{idx}"
+        outputs = []
+        for i in range(len(operation.split)):
+            outputs.append(f"output_{i}")
+        outputs = np.array(outputs)
+        x = self._to_onnx_proto(operation.x, f"{opname}.x")
+        split = self._to_onnx_proto(operation.split, f"{opname}.split")
+        node = onnx.helper.make_node(
+            op_type,
+            inputs=[x.name, split.name],
+            outputs=outputs,
+            name=opname,
+            axis=operation.axis,
+        )
+
+        return node
+
     def visit_Slice(self, operation: operations.Slice) -> onnx.NodeProto:
         op_type = str(operation)
         idx = self.op_counts[op_type] = self.op_counts[op_type] + 1

dnnv/nn/converters/tensorflow.py

@@ -241,7 +241,6 @@ def conv_func(*inputs):
             else:
                 bias = np.zeros((weights.shape[0],), dtype=weights.dtype)
             assert np.all(operation.dilations == 1)
-            assert np.all(operation.group == 1)
             num_pads = len(operation.pads)
             pads = tuple(
                 zip(
@@ -308,7 +307,6 @@ def convtranspose_func(*inputs):
                 bias = operation.b
             else:
                 bias = np.zeros((weights.shape[1],), dtype=weights.dtype)
-            assert np.all(operation.group == 1)
 
             num_pads = len(operation.pads)
             pads = tuple(
@@ -858,6 +856,23 @@ def softmax_func(*inputs):
 
         return softmax_func
 
+    def visit_Split(self, operation):
+        x_ = operation.x
+        if isinstance(x_, Operation):
+            x_ = self.visit(x_)
+        split_ = operation.split
+        if isinstance(split_, Operation):
+            split_ = self.visit(split_)
+        axis = operation.axis
+
+        @self._cached
+        def split_func(*inputs):
+            x, split = _concretize([x_, split_], inputs)
+            x = tf.split(x, tf.convert_to_tensor(split), axis=axis)
+            return x
+
+        return split_func
+
     def visit_Sub(self, operation):
         a_ = operation.a
         if isinstance(a_, Operation):

dnnv/nn/operations/tensor.py

@@ -176,6 +176,24 @@ def from_onnx(cls, onnx_node, *inputs):
         return cls(*inputs, name=onnx_node.name)
 
 
+class Split(Operation):
+    def __init__(self, x, split=None, *, axis=0, name: Optional[str] = None):
+        super().__init__(name=name)
+        self.x = x
+        self.axis = axis
+        self.split = split
+
+    @classmethod
+    def from_onnx(cls, onnx_node, *inputs):
+        attributes = {a.name: as_numpy(a) for a in onnx_node.attribute}
+        axis = attributes.get("axis", 0)
+        if len(inputs) < 2:
+            # TODO: split is an input past version 11 (?)
+            split = attributes.get("split")
+            return cls(*inputs, axis=axis, split=split, name=onnx_node.name)
+        return cls(*inputs, axis=axis, name=onnx_node.name)
+
+
 class Slice(Operation):
     def __init__(
         self, x, starts, ends, axes=None, steps=None, *, name: Optional[str] = None
@@ -249,6 +267,7 @@ def from_onnx(cls, onnx_node, *inputs):
     "Reshape",
     "Resize",
     "Shape",
+    "Split",
     "Slice",
     "Tile",
     "Transpose",

dnnv/nn/visitors.py

@@ -349,6 +349,14 @@ def visit_Softmax(self, operation: operations.Softmax) -> None:
         self.print_op_id(operation)
         print(f"Softmax({self.get_op_id(operation.x)}, axis={operation.axis})")
 
+    def visit_Split(self, operation: operations.Sub) -> None:
+        self.generic_visit(operation)
+        self.print_op_id(operation)
+        print(
+            "Split(%s, axis=%s, split=%s)"
+            % (self.get_op_id(operation.x), operation.axis, operation.split)
+        )
+
     def visit_Sub(self, operation: operations.Sub) -> None:
         self.generic_visit(operation)
         self.print_op_id(operation)

tests/unit_tests/test_nn/test_converters/test_onnx/test_Split.py

@@ -0,0 +1,113 @@
+import numpy as np
+import onnxruntime.backend
+import pytest
+
+from dnnv.nn.converters.onnx import *
+from dnnv.nn.operations import *
+
+
+def test_Split_export_1d() -> None:
+    input = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).astype(np.float32)
+    op = Split(input, axis=0, split=np.array([2, 2, 2]))
+    all_results = []
+    for i in range(3):
+        onnx_model = convert(OperationGraph([OutputSelect(op, i)]))
+        results = onnxruntime.backend.run(onnx_model, [])
+        all_results.append(results[0])
+    all_results = np.array(all_results)
+    expected_outputs = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]).astype(np.float32)
+    assert len(all_results) == 3
+    assert np.array_equiv(all_results, expected_outputs)
+
+    op = Split(input, axis=0, split=np.array([2, 4]).astype(np.int64))
+    all_results = []
+    for i in range(2):
+        onnx_model = convert(OperationGraph([OutputSelect(op, i)]))
+        results = onnxruntime.backend.run(onnx_model, [])
+        all_results.append(results[0])
+    all_results = np.array(all_results)
+    assert len(all_results) == 2
+    assert np.array_equiv(all_results[0], np.array([1.0, 2.0]).astype(np.float32))
+    assert np.array_equiv(
+        all_results[1], np.array([3.0, 4.0, 5.0, 6.0]).astype(np.float32)
+    )
+
+
+def test_Split_export_2d() -> None:
+    input = np.array(
+        [[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [7.0, 8.0, 9.0, 10.0, 11.0, 12.0]]
+    ).astype(np.float32)
+
+    op = Split(input, axis=1, split=np.array([3, 3]))
+    all_results = []
+    for i in range(2):
+        outputselect = OutputSelect(op, i)
+        onnx_model = convert(OperationGraph([outputselect]))
+        results = onnxruntime.backend.run(onnx_model, [])
+        all_results.append(results[0])
+    expected_outputs = np.array(
+        [[[1.0, 2.0, 3.0], [7.0, 8.0, 9.0]], [[4.0, 5.0, 6.0], [10.0, 11.0, 12.0]]]
+    ).astype(np.float32)
+    for i in range(2):
+        assert all_results[i].shape == (2, 3)
+        assert np.array_equiv(all_results[i], expected_outputs[i])
+
+    op = Split(input, axis=1, split=np.array([2, 4]))
+    all_results = []
+    for i in range(2):
+        outputselect = OutputSelect(op, i)
+        onnx_model = convert(OperationGraph([outputselect]))
+        results = onnxruntime.backend.run(onnx_model, [])
+        all_results.append(results[0])
+    expected_outputs1 = np.array([[1.0, 2.0], [7.0, 8.0]]).astype(np.float32)
+    expected_outputs2 = np.array(
+        [[3.0, 4.0, 5.0, 6.0], [9.0, 10.0, 11.0, 12.0]]
+    ).astype(np.float32)
+    assert all_results[0].shape == (2, 2)
+    assert np.array_equiv(all_results[0], expected_outputs1)
+    assert all_results[1].shape == (2, 4)
+    assert np.array_equiv(all_results[1], expected_outputs2)
+
+
+def test_Split_export_default_values() -> None:
+    input = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).astype(np.float32)
+    op = Split(input, split=np.array([2, 2, 2]))
+    all_results = []
+    for i in range(3):
+        onnx_model = convert(OperationGraph([OutputSelect(op, i)]))
+        results = onnxruntime.backend.run(onnx_model, [])
+        all_results.append(results[0])
+    all_results = np.array(all_results)
+    expected_outputs = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]).astype(np.float32)
+    assert len(all_results) == 3
+    assert all_results.shape == expected_outputs.shape
+    assert np.array_equiv(all_results, expected_outputs)
+
+    op = Split(input, split=np.array([2, 4]))
+    all_results = []
+    for i in range(2):
+        onnx_model = convert(OperationGraph([OutputSelect(op, i)]))
+        results = onnxruntime.backend.run(onnx_model, [])
+        all_results.append(results[0])
+    assert len(all_results) == 2
+    assert len(all_results[0]) == 2
+    assert len(all_results[1]) == 4
+    assert np.array_equiv(all_results[0], np.array([1.0, 2.0]).astype(np.float32))
+    assert np.array_equiv(
+        all_results[1], np.array([3.0, 4.0, 5.0, 6.0]).astype(np.float32)
+    )
+
+
+def test_Split_export_zero_size_splits() -> None:
+    # Split emtpy tensor to tensors of size zero
+    input = np.array([]).astype(np.float32)
+    op = Split(input, split=np.array([0, 0, 0]))
+    all_results = []
+    for i in range(3):
+        onnx_model = convert(OperationGraph([OutputSelect(op, i)]))
+        results = onnxruntime.backend.run(onnx_model, [])
+        all_results.append(results[0])
+    all_results = np.array(all_results)
+    expected_outputs = np.array([[], [], []]).astype(np.float32)
+    assert all_results.shape == (3, 0)
+    assert np.array_equiv(all_results, expected_outputs)

tests/unit_tests/test_nn/test_converters/test_tensorflow/test_Split.py

@@ -0,0 +1,80 @@
+import numpy as np
+import pytest
+
+from dnnv.nn.converters.tensorflow import *
+from dnnv.nn.operations import *
+
+
+def test_Split_export_1d() -> None:
+    input = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).astype(np.float32)
+    op = Split(input, axis=0, split=np.array([2, 2, 2]))
+    tf_op = TensorflowConverter().visit(op)
+    result_ = tf_op()
+    expected_outputs = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]).astype(np.float32)
+    assert len(result_) == 3
+    assert np.array_equiv(result_, expected_outputs)
+
+    op = Split(input, axis=0, split=np.array([2, 4]).astype(np.int64))
+    tf_op = TensorflowConverter().visit(op)
+    result_ = tf_op()
+    assert len(result_) == 2
+    assert np.array_equiv(result_[0], np.array([1.0, 2.0]).astype(np.float32))
+    assert np.array_equiv(result_[1], np.array([3.0, 4.0, 5.0, 6.0]).astype(np.float32))
+
+
+def test_Split_export_2d() -> None:
+    input = np.array(
+        [[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [7.0, 8.0, 9.0, 10.0, 11.0, 12.0]]
+    ).astype(np.float32)
+
+    op = Split(input, axis=1, split=np.array([3, 3]))
+    tf_op = TensorflowConverter().visit(op)
+    result_ = tf_op()
+    expected_outputs = np.array(
+        [[[1.0, 2.0, 3.0], [7.0, 8.0, 9.0]], [[4.0, 5.0, 6.0], [10.0, 11.0, 12.0]]]
+    ).astype(np.float32)
+    for i in range(2):
+        assert result_[i].shape == (2, 3)
+        assert np.array_equiv(result_[i], expected_outputs[i])
+
+    op = Split(input, axis=1, split=np.array([2, 4]))
+    tf_op = TensorflowConverter().visit(op)
+    result_ = tf_op()
+    expected_outputs1 = np.array([[1.0, 2.0], [7.0, 8.0]]).astype(np.float32)
+    expected_outputs2 = np.array(
+        [[3.0, 4.0, 5.0, 6.0], [9.0, 10.0, 11.0, 12.0]]
+    ).astype(np.float32)
+    assert result_[0].shape == (2, 2)
+    assert np.array_equiv(result_[0], expected_outputs1)
+    assert result_[1].shape == (2, 4)
+    assert np.array_equiv(result_[1], expected_outputs2)
+
+
+def test_Split_export_default_values() -> None:
+    input = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).astype(np.float32)
+    op = Split(input, split=np.array([2, 2, 2]))
+    tf_op = TensorflowConverter().visit(op)
+    result_ = tf_op()
+    expected_outputs = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]).astype(np.float32)
+    assert len(result_) == 3
+    assert np.array(result_).shape == expected_outputs.shape
+    assert np.array_equiv(np.array(result_), expected_outputs)
+
+    op = Split(input, split=np.array([2, 4]))
+    tf_op = TensorflowConverter().visit(op)
+    result_ = tf_op()
+    assert len(result_) == 2
+    assert len(result_[0]) == 2
+    assert len(result_[1]) == 4
+    assert np.array_equiv(result_[0], np.array([1.0, 2.0]).astype(np.float32))
+    assert np.array_equiv(result_[1], np.array([3.0, 4.0, 5.0, 6.0]).astype(np.float32))
+
+
+def test_Split_export_zero_size_splits() -> None:
+    input = np.array([]).astype(np.float32)
+    op = Split(input, split=np.array([0, 0, 0]))
+    tf_op = TensorflowConverter().visit(op)
+    result_ = tf_op()
+    expected_outputs = np.array([[], [], []]).astype(np.float32)
+    assert np.array(result_).shape == (3, 0)
+    assert np.array_equiv(np.array(result_), expected_outputs)