Add ONNX SpaceToDepth and DepthToSpace extractor (#1122)

* Add ONNX SpaceToDepth extractor

* Add DepthToSpace ONNX extractor

diff --git a/model-optimizer/automation/package_BOM.txt b/model-optimizer/automation/package_BOM.txt
index 6b47d9f38b6b5fb9e2040734a81c61f609a1eb2e..b4a08df99da0a7630c211e528e470fc16db06fb2 100644 (file)
--- a/model-optimizer/automation/package_BOM.txt
+++ b/model-optimizer/automation/package_BOM.txt
@@ -241,6 +241,7 @@ extensions/front/onnx/conv_ext.py
 extensions/front/onnx/crop_ext.py
 extensions/front/onnx/cumsum_ext.py
 extensions/front/onnx/deformable_conv_ext.py
+extensions/front/onnx/depth_to_space_ext.py
 extensions/front/onnx/detection_output.py
 extensions/front/onnx/detectionoutput_ext.py
 extensions/front/onnx/dropout_ext.py
@@ -302,6 +303,7 @@ extensions/front/onnx/slice_ext.py
 extensions/front/onnx/softmax_ext.py
 extensions/front/onnx/softmaxONNX_to_softmax.py
 extensions/front/onnx/softplus_ext.py
+extensions/front/onnx/space_to_depth_ext.py
 extensions/front/onnx/split_ext.py
 extensions/front/onnx/squeeze_ext.py
 extensions/front/onnx/top_k_ext.py

diff --git a/model-optimizer/extensions/front/onnx/depth_to_space_ext.py b/model-optimizer/extensions/front/onnx/depth_to_space_ext.py
new file mode 100644 (file)
index 0000000..f043a1c
--- /dev/null
+++ b/model-optimizer/extensions/front/onnx/depth_to_space_ext.py
@@ -0,0 +1,42 @@
+"""
+ Copyright (C) 2020 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+import logging as log
+
+from extensions.ops.depth_to_space import DepthToSpaceOp
+from mo.front.extractor import FrontExtractorOp
+from mo.front.onnx.extractors.utils import onnx_attr
+
+
+class DepthToSpaceFrontExtractor(FrontExtractorOp):
+    op = 'DepthToSpace'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        # update the attributes of the node
+        node_name = node.soft_get('name', node.id)
+        block_size = onnx_attr(node, 'blocksize', 'i', default=None)
+        assert block_size is not None, \
+            'DepthToSpace should have "blocksize" attribute specified for node {}'.format(node_name)
+        onnx_mode = onnx_attr(node, 'mode', 's', default=b'DCR').decode()
+        assert onnx_mode in ['DCR', 'CRD'], 'Unrecognized mode provided for DepthToSpace node {}'.format(node_name)
+        if onnx_mode == 'DCR':
+            mode = 'depth_first'
+        else:
+            mode = 'blocks_first'
+
+        DepthToSpaceOp.update_node_stat(node, {'block_size': block_size, 'mode': mode})
+        return cls.enabled

diff --git a/model-optimizer/extensions/front/onnx/space_to_depth_ext.py b/model-optimizer/extensions/front/onnx/space_to_depth_ext.py
new file mode 100644 (file)
index 0000000..c468dd2
--- /dev/null
+++ b/model-optimizer/extensions/front/onnx/space_to_depth_ext.py
@@ -0,0 +1,31 @@
+"""
+ Copyright (C) 2020 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+
+from extensions.ops.space_to_depth import SpaceToDepth
+from mo.front.extractor import FrontExtractorOp
+from mo.front.onnx.extractors.utils import onnx_attr
+
+
+class SpaceToDepthFrontExtractor(FrontExtractorOp):
+    op = 'SpaceToDepth'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        # update the attributes of the node
+        block_size = onnx_attr(node, 'blocksize', 'i', default=None)
+        SpaceToDepth.update_node_stat(node, {'block_size': block_size})
+        return cls.enabled