From 041871608b613d8887ba0cf0693f97ff035428ba Mon Sep 17 00:00:00 2001 From: Evgeny Lazarev Date: Thu, 20 Aug 2020 13:57:22 +0300 Subject: [PATCH] ONNX ScatterND MO support (#1813) * Added extractor for ONNX operation ScatterNDUpdate * Updated list of supported operations for ONNX --- .../prepare_model/Supported_Frameworks_Layers.md | 1 + model-optimizer/extensions/front/onnx/scatter_ext.py | 11 +++++++++++ 2 files changed, 12 insertions(+) diff --git a/docs/MO_DG/prepare_model/Supported_Frameworks_Layers.md b/docs/MO_DG/prepare_model/Supported_Frameworks_Layers.md index 7cf9492d5edc4a..0f9bc8da7e738e 100644 --- a/docs/MO_DG/prepare_model/Supported_Frameworks_Layers.md +++ b/docs/MO_DG/prepare_model/Supported_Frameworks_Layers.md @@ -379,6 +379,7 @@ Standard ONNX\* operators: | Resize | Opset-10 version is supported | | ReverseSequence | No | | Scatter | Supported if fuse-able to ScatterUpdate. MYRIAD only | +| ScatterND | No | | ScatterElements | Supported if fuse-able to ScatterUpdate. MYRIAD only | | Select | No | | Shape | No | diff --git a/model-optimizer/extensions/front/onnx/scatter_ext.py b/model-optimizer/extensions/front/onnx/scatter_ext.py index 8d3be1c827c3d2..8e9a1d2bd86e26 100644 --- a/model-optimizer/extensions/front/onnx/scatter_ext.py +++ b/model-optimizer/extensions/front/onnx/scatter_ext.py @@ -14,6 +14,7 @@ limitations under the License. """ from extensions.ops.scatter import ScatterElementsUpdate +from extensions.ops.scatternd import ScatterNDUpdate from mo.front.extractor import FrontExtractorOp from mo.front.onnx.extractors.utils import onnx_attr @@ -39,3 +40,13 @@ def extract(cls, node): axis = onnx_attr(node, 'axis', 'i', default=0) ScatterElementsUpdate.update_node_stat(node, {'axis': axis}) return cls.enabled + + +class ScatterNDExtractor(FrontExtractorOp): + op = 'ScatterND' + enabled = True + + @classmethod + def extract(cls, node): + ScatterNDUpdate.update_node_stat(node, {}) + return cls.enabled