ONNX ScatterND MO support (openvinotoolkit#1813)

* Added extractor for ONNX operation ScatterNDUpdate * Updated list of supported operations for ONNX
RomanZm · Aug 28, 2020 · 0418716 · 0418716
1 parent d800a0a
commit 0418716
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git 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
@@ -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