From 34d74d6bfc7de5fba0cce2b9eb52d5ba64f27973 Mon Sep 17 00:00:00 2001 From: Guriido Date: Thu, 28 Sep 2023 08:27:21 +0900 Subject: [PATCH 1/4] allow any split --- nnoir-onnx/README.md | 2 - nnoir-onnx/nnoir_onnx/operators/split.py | 70 +++++++++++------------- nnoir-onnx/test/test_split.py | 37 +++++++++++-- 3 files changed, 62 insertions(+), 47 deletions(-) diff --git a/nnoir-onnx/README.md b/nnoir-onnx/README.md index 5fcc36d..776e3cb 100644 --- a/nnoir-onnx/README.md +++ b/nnoir-onnx/README.md @@ -91,8 +91,6 @@ docker run --rm -it -u $UID:$GID -v $(pwd):/work idein/nnoir-tools:20230720 onnx * [Sin](https://github.com/onnx/onnx/blob/master/docs/Operators.md#Sin) * [Softmax](https://github.com/onnx/onnx/blob/master/docs/Operators.md#Softmax) * [Split](https://github.com/onnx/onnx/blob/master/docs/Operators.md#Split) - * must be from opset version >= 13 - * Second optional parameter `split` is not supported * [Squeeze](https://github.com/onnx/onnx/blob/master/docs/Operators.md#Squeeze) * [Sub](https://github.com/onnx/onnx/blob/master/docs/Operators.md#Sub) * 1st input must not be `"constant"` diff --git a/nnoir-onnx/nnoir_onnx/operators/split.py b/nnoir-onnx/nnoir_onnx/operators/split.py index 0c42117..c3a080d 100644 --- a/nnoir-onnx/nnoir_onnx/operators/split.py +++ b/nnoir-onnx/nnoir_onnx/operators/split.py @@ -1,4 +1,4 @@ -from typing import Tuple +from typing import List import numpy as np from nnoir.functions import Constant, MatMul, Transpose @@ -6,13 +6,19 @@ from .utils import Op, UnsupportedONNXOperation, gen_unregisterd_node_name, register_node -def create_half_split_matrices(k: int) -> Tuple[np.ndarray, np.ndarray]: - k_2 = k // 2 +def create_split_matrices(k: int, sizes: List[int]) -> List[np.ndarray]: - eye = np.eye(k_2, dtype="float32") - zero = np.zeros((k_2, k_2), dtype="float32") + split_matrices = [] + acc = 0 + for sub_k in sizes: + zero_before = np.zeros((acc, sub_k), dtype="float32") + eye = np.eye(sub_k, dtype="float32") + zero_after = np.zeros((k - sub_k - acc, sub_k), dtype="float32") - return (np.concatenate([eye, zero]), np.concatenate([zero, eye])) + split_matrices.append(np.concatenate([zero_before, eye, zero_after])) + acc += sub_k + + return split_matrices def gen_value(env, arr): @@ -35,22 +41,23 @@ def __init__(self, node, *args): super().__init__(node, *args) def to_function(self, env, constants): - if len(self.node.input) > 1: - raise UnsupportedONNXOperation(self.node, "the number of inputs must be 1.") - if len(self.node.output) > 2: - raise UnsupportedONNXOperation(self.node, "the number of outputs must be 2.") - split_axis = 0 for attr in self.node.attribute: if attr.name == "axis": split_axis = attr.i + output_sizes_on_axis = [] + for output in self.node.output: + output_sizes_on_axis.append(env[output].shape[split_axis]) + shape = env[self.node.input[0]].shape k = shape[split_axis] - if k % 2 != 0: - raise Exception("Cannot reshape on odd size, shape {}, axis {}".format(shape, split_axis)) + assert k == sum( + output_sizes_on_axis + ), f"outputs are not a split of input on for axis {split_axis} of input shape {shape}: {output_sizes_on_axis}" + + matrices = create_split_matrices(k, output_sizes_on_axis) - matrice_up, matrice_down = create_half_split_matrices(k) transpose_perm_0 = list(range(len(shape))) transpose_perm_0.append(transpose_perm_0.pop(split_axis)) transpose_perm_1 = list(range(len(shape))) @@ -77,30 +84,15 @@ def linear_shape(x_shape, w_shape): trans_shape = transpose_shape(shape, transpose_perm_0) trans_out = gen_dummy_value(env, trans_shape) - linear_up_shape = linear_shape(trans_shape, matrice_up.shape) - linear_up_out = gen_dummy_value(env, linear_up_shape) - - linear_down_shape = linear_shape(trans_shape, matrice_down.shape) - linear_down_out = gen_dummy_value(env, linear_down_shape) - - up_const = gen_value(env, matrice_up) - up_const_node = Constant([], [up_const], value=matrice_up) - down_const = gen_value(env, matrice_up) - down_const_node = Constant([], [down_const], value=matrice_down) - - transpose_node = Transpose(list(self.node.input), [trans_out], axes=transpose_perm_0) - linear_up_node = MatMul([trans_out, up_const], [linear_up_out]) - linear_down_node = MatMul([trans_out, down_const], [linear_down_out]) - transpose_up_node = Transpose([linear_up_out], [self.node.output[0]], axes=transpose_perm_1) - transpose_down_node = Transpose([linear_down_out], [self.node.output[1]], axes=transpose_perm_1) - nodes = [ - up_const_node, - down_const_node, - transpose_node, - linear_up_node, - linear_down_node, - transpose_up_node, - transpose_down_node, - ] + nodes = [Transpose(list(self.node.input), [trans_out], axes=transpose_perm_0)] + for i, mat in enumerate(matrices): + _linear_shape = linear_shape(trans_shape, mat.shape) + linear_out = gen_dummy_value(env, _linear_shape) + + _const = gen_value(env, mat) + _const_node = Constant([], [_const], value=mat) + linear_node = MatMul([trans_out, _const], [linear_out]) + transpose_node = Transpose([linear_out], [self.node.output[i]], axes=transpose_perm_1) + nodes.extend([_const_node, linear_node, transpose_node]) return nodes diff --git a/nnoir-onnx/test/test_split.py b/nnoir-onnx/test/test_split.py index ff58b69..eb826c1 100644 --- a/nnoir-onnx/test/test_split.py +++ b/nnoir-onnx/test/test_split.py @@ -83,11 +83,37 @@ def create_onnx(self) -> onnx.ModelProto: SplitTester({"v0": v0}, outputs).run() -@pytest.mark.xfail() def test_split_specify_split(): """ - Specify second input (optional parameter). - Due to lack of implementation, the second input is not supported. + Specify split attribute (opset 11). + """ + + class SplitTester(Base): + def __init__(self, inputs, outputs): + super().__init__(inputs, outputs) + + def create_onnx(self) -> onnx.ModelProto: + node = make_node("Split", inputs=["v0"], outputs=["v1", "v2", "v3"], axis=3, split=[2, 3, 5]) + inputs = [ + info("v0", TensorProto.FLOAT, (1, 3, 4, 10)), + ] + outputs = [ + info("v1", TensorProto.FLOAT, (1, 3, 4, 2)), + info("v2", TensorProto.FLOAT, (1, 3, 4, 3)), + info("v3", TensorProto.FLOAT, (1, 3, 4, 5)), + ] + + graph = make_graph([node], "add_graph", inputs, outputs) + return make_model(graph, opset_imports=[make_opsetid("", 11)]) + + v0 = np.random.rand(1, 3, 4, 10).astype(np.float32) + outputs = ["v1", "v2", "v3"] + SplitTester({"v0": v0}, outputs).run() + + +def test_split_specify_split_13(): + """ + Specify split input (opset 13). """ class SplitTester(Base): @@ -107,11 +133,10 @@ def create_onnx(self) -> onnx.ModelProto: ] graph = make_graph([node], "add_graph", inputs, outputs) - model = make_model(graph) - return model + return make_model(graph, opset_imports=[make_opsetid("", 13)]) v0 = np.random.rand(1, 3, 4, 10).astype(np.float32) p0 = np.array([2, 3, 5]).astype(np.int64) - outputs = ["v1", "v2"] + outputs = ["v1", "v2", "v3"] SplitTester({"v0": v0, "p0": p0}, outputs).run() From ce560c45ef82c5a1f8c129843ea382a8fe693ab4 Mon Sep 17 00:00:00 2001 From: Guriido Date: Thu, 28 Sep 2023 10:08:46 +0900 Subject: [PATCH 2/4] fix split --- nnoir-onnx/nnoir_onnx/operators/split.py | 2 +- nnoir-onnx/test/test_split.py | 14 +++++++++----- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/nnoir-onnx/nnoir_onnx/operators/split.py b/nnoir-onnx/nnoir_onnx/operators/split.py index c3a080d..c060cf7 100644 --- a/nnoir-onnx/nnoir_onnx/operators/split.py +++ b/nnoir-onnx/nnoir_onnx/operators/split.py @@ -84,7 +84,7 @@ def linear_shape(x_shape, w_shape): trans_shape = transpose_shape(shape, transpose_perm_0) trans_out = gen_dummy_value(env, trans_shape) - nodes = [Transpose(list(self.node.input), [trans_out], axes=transpose_perm_0)] + nodes = [Transpose([self.node.input[0]], [trans_out], axes=transpose_perm_0)] for i, mat in enumerate(matrices): _linear_shape = linear_shape(trans_shape, mat.shape) linear_out = gen_dummy_value(env, _linear_shape) diff --git a/nnoir-onnx/test/test_split.py b/nnoir-onnx/test/test_split.py index eb826c1..a05bb6d 100644 --- a/nnoir-onnx/test/test_split.py +++ b/nnoir-onnx/test/test_split.py @@ -2,7 +2,7 @@ import onnx import pytest from onnx import TensorProto -from onnx.helper import make_graph, make_model, make_node, make_opsetid, make_tensor_value_info +from onnx.helper import make_graph, make_model, make_node, make_opsetid, make_tensor_value_info, make_tensor from onnx.numpy_helper import from_array from util import Base @@ -124,19 +124,23 @@ def create_onnx(self) -> onnx.ModelProto: node = make_node("Split", inputs=["v0", "p0"], outputs=["v1", "v2", "v3"], axis=3) inputs = [ info("v0", TensorProto.FLOAT, (1, 3, 4, 10)), - info("p0", TensorProto.INT64, (3,)), ] + node_p0 = make_node( + "Constant", + value=make_tensor(name="p0_constant", data_type=TensorProto.INT64, dims=(3,), vals=np.array([2, 3, 5]).astype(np.int64)), + inputs=[], + outputs=["p0"], + ) outputs = [ info("v1", TensorProto.FLOAT, (1, 3, 4, 2)), info("v2", TensorProto.FLOAT, (1, 3, 4, 3)), info("v3", TensorProto.FLOAT, (1, 3, 4, 5)), ] - graph = make_graph([node], "add_graph", inputs, outputs) + graph = make_graph([node_p0, node], "add_graph", inputs, outputs) return make_model(graph, opset_imports=[make_opsetid("", 13)]) v0 = np.random.rand(1, 3, 4, 10).astype(np.float32) - p0 = np.array([2, 3, 5]).astype(np.int64) outputs = ["v1", "v2", "v3"] - SplitTester({"v0": v0, "p0": p0}, outputs).run() + SplitTester({"v0": v0}, outputs).run() From 3e3cb88f879a43a41fd908d677adb546956df1fa Mon Sep 17 00:00:00 2001 From: Guriido Date: Thu, 28 Sep 2023 10:09:17 +0900 Subject: [PATCH 3/4] fix lint --- nnoir-onnx/test/test_split.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/nnoir-onnx/test/test_split.py b/nnoir-onnx/test/test_split.py index a05bb6d..45ba7b8 100644 --- a/nnoir-onnx/test/test_split.py +++ b/nnoir-onnx/test/test_split.py @@ -2,7 +2,7 @@ import onnx import pytest from onnx import TensorProto -from onnx.helper import make_graph, make_model, make_node, make_opsetid, make_tensor_value_info, make_tensor +from onnx.helper import make_graph, make_model, make_node, make_opsetid, make_tensor, make_tensor_value_info from onnx.numpy_helper import from_array from util import Base @@ -127,7 +127,9 @@ def create_onnx(self) -> onnx.ModelProto: ] node_p0 = make_node( "Constant", - value=make_tensor(name="p0_constant", data_type=TensorProto.INT64, dims=(3,), vals=np.array([2, 3, 5]).astype(np.int64)), + value=make_tensor( + name="p0_constant", data_type=TensorProto.INT64, dims=(3,), vals=np.array([2, 3, 5]).astype(np.int64) + ), inputs=[], outputs=["p0"], ) From 10b67531fa25b5a243cc9d537bbc246499f119b2 Mon Sep 17 00:00:00 2001 From: Guriido Date: Mon, 2 Oct 2023 09:05:13 +0900 Subject: [PATCH 4/4] apply linter --- nnoir-onnx/nnoir_onnx/operators/split.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/nnoir-onnx/nnoir_onnx/operators/split.py b/nnoir-onnx/nnoir_onnx/operators/split.py index a7bddaa..faf0bd0 100644 --- a/nnoir-onnx/nnoir_onnx/operators/split.py +++ b/nnoir-onnx/nnoir_onnx/operators/split.py @@ -9,6 +9,7 @@ ShapeLike = Union[Tuple[int, ...], List[int]] + def create_split_matrices(k: int, sizes: List[int]) -> List[NDArray[Any]]: split_matrices: List[NDArray[Any]] = [] acc = 0 @@ -96,5 +97,5 @@ def linear_shape(x_shape: ShapeLike, w_shape: ShapeLike) -> ShapeLike: linear_node = MatMul([trans_out, _const], [linear_out]) # type: ignore transpose_node = Transpose([linear_out], [self.node.output[i]], axes=transpose_perm_1) # type: ignore nodes.extend([_const_node, linear_node, transpose_node]) - + return nodes