Add ONNX SpaceToDepth and DepthToSpace extractor (#1122)

* Add ONNX SpaceToDepth extractor

* Add DepthToSpace ONNX extractor
This commit is contained in:
Maxim Vafin 2020-07-23 13:05:42 +03:00 committed by GitHub
parent f00cdddede
commit 0063efeb09
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 75 additions and 0 deletions

View File

@ -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

View File

@ -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

View File

@ -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