From 538f407829313b3e4340d75a310ed09c0bd7cb58 Mon Sep 17 00:00:00 2001 From: Anton Chetverikov Date: Wed, 19 May 2021 19:03:58 +0300 Subject: [PATCH] Fix errors for VariadicSplit layer restored from serialized IR (#5613) * Fix errors in VariadicSplit layer restored from serialized IR * Update VariadicSplit specification and error message to allow 1D tensors on 1st input * Update spec * Resolve comments * Apply comments, add unit tests * Update unit tests --- docs/ops/movement/VariadicSplit_1.md | 4 +- model-optimizer/extensions/ops/split.py | 3 +- .../mo/utils/ir_reader/layer_to_class.py | 3 +- .../unit_tests/extensions/ops/split_test.py | 82 ++++++++++++++++++- 4 files changed, 87 insertions(+), 5 deletions(-) diff --git a/docs/ops/movement/VariadicSplit_1.md b/docs/ops/movement/VariadicSplit_1.md index 87c9ac67b65..9b2392d6fc9 100644 --- a/docs/ops/movement/VariadicSplit_1.md +++ b/docs/ops/movement/VariadicSplit_1.md @@ -8,7 +8,7 @@ **Detailed Description** -*VariadicSplit* operation splits a given input tensor `data` into chunks along a scalar `axis`. It produces multiple output tensors based on additional input tensor `split_lengths`. +*VariadicSplit* operation splits a given input tensor `data` into chunks along a scalar or tensor with shape `[1]` `axis`. It produces multiple output tensors based on additional input tensor `split_lengths`. The i-th output tensor shape is equal to the input tensor `data` shape, except for dimension along `axis` which is `split_lengths[i]`. \f[ @@ -23,7 +23,7 @@ Where D is the rank of input tensor `data`. The sum of elements in `split_length * **1**: `data`. A tensor of type `T1` and arbitrary shape. **Required.** -* **2**: `axis`. Axis along `data` to split. A scalar of type `T2` with value from range `-rank(data) .. rank(data)-1`. Negative values address dimensions from the end. +* **2**: `axis`. Axis along `data` to split. A scalar or tensor with shape `[1]` of type `T2` with value from range `-rank(data) .. rank(data)-1`. Negative values address dimensions from the end. **Required.** * **3**: `split_lengths`. A list containing the dimension values of each output tensor shape along the split `axis`. A 1D tensor of type `T2`. The number of elements in `split_lengths` determines the number of outputs. The sum of elements in `split_lengths` must match `data.shape[axis]`. In addition `split_lengths` can contain a single `-1` element, which means, all remaining items along specified `axis` that are not consumed by other parts. **Required.** diff --git a/model-optimizer/extensions/ops/split.py b/model-optimizer/extensions/ops/split.py index 211adcb4d73..6c25aa7ef20 100644 --- a/model-optimizer/extensions/ops/split.py +++ b/model-optimizer/extensions/ops/split.py @@ -44,7 +44,8 @@ class VariadicSplitBase(Op): axis = node.in_port(1).data.get_value() if op == 'VariadicSplit' else node.soft_get('axis', None) assert axis is not None, '{} `axis` is unknown for node {}'.format(op, name) - assert axis.ndim == 0, '{} `axis` should be scalar, but it`s not for node {}'.format(op, name) + assert axis.ndim == 0 or (axis.ndim == 1 and axis.shape[0] == 1), \ + '{} `axis` should be scalar or tensor with shape [1], but it`s not for node {}'.format(op, name) split_lengths = node.in_port(2).data.get_value() if op == 'VariadicSplit' else node.soft_get('split_lengths', None) diff --git a/model-optimizer/mo/utils/ir_reader/layer_to_class.py b/model-optimizer/mo/utils/ir_reader/layer_to_class.py index 90f07c6c5bf..d1afed6918d 100644 --- a/model-optimizer/mo/utils/ir_reader/layer_to_class.py +++ b/model-optimizer/mo/utils/ir_reader/layer_to_class.py @@ -275,7 +275,8 @@ postprocessing_op_nodes = { 'TensorIterator': ti_add_edge_attrs, 'TopK': TopKNormalizer.normalize_outputs, # Call normalize Split outputs for generated IR by ir-reader - 'Split': AddFakeOutputsToSplit.split_normalize_outputs + 'Split': AddFakeOutputsToSplit.split_normalize_outputs, + 'VariadicSplit': AddFakeOutputsToSplit.split_normalize_outputs, } diff --git a/model-optimizer/unit_tests/extensions/ops/split_test.py b/model-optimizer/unit_tests/extensions/ops/split_test.py index 0ab4f40d0e3..1a1accca1b5 100644 --- a/model-optimizer/unit_tests/extensions/ops/split_test.py +++ b/model-optimizer/unit_tests/extensions/ops/split_test.py @@ -5,12 +5,14 @@ import unittest import numpy as np -from extensions.ops.split import AttributedSplit, AttributedVariadicSplit +from extensions.ops.split import AttributedSplit, AttributedVariadicSplit, VariadicSplit from mo.front.common.partial_infer.utils import int64_array from mo.graph.graph import Node from mo.utils.ir_engine.compare_graphs import compare_graphs from unit_tests.utils.graph import build_graph +from generator import generator, generate + class TestSplitOp(unittest.TestCase): nodes = { @@ -188,3 +190,81 @@ class TestAttributedVariadicSplitOp(unittest.TestCase): self.assertTrue(node.out_port(4).disconnected()) self.assertTrue(node.out_port(3).disconnected()) self.assertTrue(np.all(node.split_lengths == np.array([2, 13, 10]))) + + +@generator +class TestVariadicSplitOp(unittest.TestCase): + nodes = { + 'input': {'kind': 'op'}, + 'split_input_data': {'kind': 'data', 'shape': None, 'value': None}, + 'split_axis': {'kind': 'op', 'op': 'Const'}, + 'split_axis_data': {'kind': 'data', 'shape': None, 'value': None}, + 'split_lengths': {'kind': 'op', 'op': 'Const'}, + 'split_lengths_data': {'kind': 'data', 'shape': None, 'value': None}, + 'split_op': {'kind': 'op', 'op': 'VariadicSplit'}, + 'split_output_0_data': {'kind': 'data', 'shape': None, 'value': None}, + 'output_0': {'kind': 'op'}, + 'split_output_1_data': {'kind': 'data', 'shape': None, 'value': None}, + 'output_1': {'kind': 'op'}, + 'split_output_2_data': {'kind': 'data', 'shape': None, 'value': None}, + 'output_2': {'kind': 'op'}, + } + edges = [ + ('input', 'split_input_data'), + ('split_input_data', 'split_op'), + ('split_axis', 'split_axis_data'), + ('split_axis_data', 'split_op'), + ('split_lengths', 'split_lengths_data'), + ('split_lengths_data', 'split_op'), + ('split_op', 'split_output_0_data'), + ('split_output_0_data', 'output_0'), + ('split_op', 'split_output_1_data'), + ('split_output_1_data', 'output_1'), + ('split_op', 'split_output_2_data'), + ('split_output_2_data', 'output_2'), + ] + + @generate(*[int64_array(2), + int64_array([2])]) + def test_variadic_split_axis(self, axis): + lengths = int64_array([2, 13, 10]) + graph = build_graph(self.nodes, self.edges, + { + 'split_input_data': {'shape': int64_array([2, 12, 25, 30])}, + 'split_axis_data': {'value': axis}, + 'split_lengths_data': {'value': lengths}, + 'split_op': {'out_ports_count': 4}, + } + ) + node = Node(graph, 'split_op') + for p in range(len(node.out_edges()), node.out_ports_count): + node.add_output_port(p) + + VariadicSplit.infer(node) + + ont_nodes_count = len(node.out_edges()) + self.assertTrue(ont_nodes_count == 3) + for out in range(ont_nodes_count): + self.assertTrue(np.all(node.out_node(out).shape == int64_array([2, 12, lengths[out], 30]))) + + @generate(*[int64_array([[2], [2]]), + int64_array([2, 2])]) + def test_negative_variadic_split_axis(self, axis): + lengths = int64_array([2, 13, 10]) + graph = build_graph(self.nodes, self.edges, + { + 'split_input_data': {'shape': int64_array([2, 12, 25, 30])}, + 'split_axis_data': {'value': axis}, + 'split_lengths_data': {'value': lengths}, + 'split_op': {'out_ports_count': 4}, + } + ) + node = Node(graph, 'split_op') + for p in range(len(node.out_edges()), node.out_ports_count): + node.add_output_port(p) + + try: + VariadicSplit.infer(node) + except AssertionError as e: + self.assertTrue(e.args[0] == 'VariadicSplit `axis` should be scalar or tensor with shape [1], ' + 'but it`s not for node split_op')