diff --git a/docs/MO_DG/prepare_model/Supported_Frameworks_Layers.md b/docs/MO_DG/prepare_model/Supported_Frameworks_Layers.md index 6c5b05a1337..03a58a23387 100644 --- a/docs/MO_DG/prepare_model/Supported_Frameworks_Layers.md +++ b/docs/MO_DG/prepare_model/Supported_Frameworks_Layers.md @@ -500,6 +500,7 @@ Standard ONNX\* operators: | Sigmoid | No | | Sign | No | | Sin | No | +| Size | No | | Slice | No | | Softmax | No | | Softplus | No | diff --git a/model-optimizer/automation/package_BOM.txt b/model-optimizer/automation/package_BOM.txt index eb216c07199..128c9c98117 100644 --- a/model-optimizer/automation/package_BOM.txt +++ b/model-optimizer/automation/package_BOM.txt @@ -338,6 +338,7 @@ extensions/front/onnx/roialign_ext.py extensions/front/onnx/roifeatureextractor_ext.py extensions/front/onnx/scatter_ext.py extensions/front/onnx/shape_ext.py +extensions/front/onnx/size_ext.py extensions/front/onnx/slice_ext.py extensions/front/onnx/softmax_ext.py extensions/front/onnx/softmaxONNX_to_softmax.py @@ -363,6 +364,7 @@ extensions/front/reshape_dim_normalizer.py extensions/front/restore_ports.py extensions/front/RollWithEmptyAxesReplacer.py extensions/front/scatter_normalizer.py +extensions/front/SizeReplacer.py extensions/front/softmax.py extensions/front/Softplus_fusion.py extensions/front/softsign_replacer.py @@ -483,7 +485,6 @@ extensions/front/tf/roll_ext.py extensions/front/tf/RollRealImagPack.py extensions/front/tf/select_ext.py extensions/front/tf/sign_ext.py -extensions/front/tf/SizeReplacer.py extensions/front/tf/slice_ext.py extensions/front/tf/softmax_ext.py extensions/front/tf/softplus_ext.py diff --git a/model-optimizer/extensions/front/tf/SizeReplacer.py b/model-optimizer/extensions/front/SizeReplacer.py similarity index 100% rename from model-optimizer/extensions/front/tf/SizeReplacer.py rename to model-optimizer/extensions/front/SizeReplacer.py diff --git a/model-optimizer/extensions/front/onnx/size_ext.py b/model-optimizer/extensions/front/onnx/size_ext.py new file mode 100644 index 00000000000..b7dba6ee28a --- /dev/null +++ b/model-optimizer/extensions/front/onnx/size_ext.py @@ -0,0 +1,17 @@ +# Copyright (C) 2018-2021 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np + +from extensions.ops.size import Size +from mo.front.extractor import FrontExtractorOp + + +class SizeExtractor(FrontExtractorOp): + op = 'Size' + enabled = True + + @classmethod + def extract(cls, node): + Size.update_node_stat(node, {'output_type': np.int64}) + return cls.enabled diff --git a/model-optimizer/extensions/ops/size.py b/model-optimizer/extensions/ops/size.py index cbfbb341b6f..a1bd0db210f 100644 --- a/model-optimizer/extensions/ops/size.py +++ b/model-optimizer/extensions/ops/size.py @@ -9,6 +9,7 @@ from mo.ops.op import Op class Size(Op): op = 'Size' + enabled = False def __init__(self, graph: Graph, attrs: dict): assert 'output_type' in attrs, 'Size has mandatory `output_type` attribute' diff --git a/model-optimizer/unit_tests/extensions/front/tf/size_replacer_test.py b/model-optimizer/unit_tests/extensions/front/size_replacer_test.py similarity index 97% rename from model-optimizer/unit_tests/extensions/front/tf/size_replacer_test.py rename to model-optimizer/unit_tests/extensions/front/size_replacer_test.py index c6f075b7084..66712c62ed2 100644 --- a/model-optimizer/unit_tests/extensions/front/tf/size_replacer_test.py +++ b/model-optimizer/unit_tests/extensions/front/size_replacer_test.py @@ -6,7 +6,7 @@ import unittest import numpy as np from generator import generator, generate -from extensions.front.tf.SizeReplacer import SizeFrontReplacer +from extensions.front.SizeReplacer import SizeFrontReplacer from mo.front.common.partial_infer.utils import int64_array from mo.utils.ir_engine.compare_graphs import compare_graphs from unit_tests.utils.graph import build_graph, regular_op_with_empty_data, result, connect, \