zhreshold closed pull request #12731: [MXNET-892] ONNX export/import:
DepthToSpace, SpaceToDepth operators
URL: https://github.com/apache/incubator-mxnet/pull/12731
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
index 204195d7409..20cfe072f9f 100644
--- a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
+++ b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
@@ -2077,3 +2077,53 @@ def convert_sqrt(node, **kwargs):
name=name,
)
return [node]
+
+@mx_op.register("depth_to_space")
+def convert_depthtospace(node, **kwargs):
+ """Map MXNet's depth_to_space operator attributes to onnx's
+ DepthToSpace operator and return the created node.
+ """
+ onnx = import_onnx_modules()
+ name = node["name"]
+ proc_nodes = kwargs["proc_nodes"]
+ inputs = node["inputs"]
+ attrs = node["attrs"]
+
+ input_node_id = kwargs["index_lookup"][inputs[0][0]]
+ input_node = proc_nodes[input_node_id].name
+
+ blksize = int(attrs.get("block_size", 0))
+
+ node = onnx.helper.make_node(
+ "DepthToSpace",
+ [input_node],
+ [name],
+ blocksize=blksize,
+ name=name,
+ )
+ return [node]
+
+@mx_op.register("space_to_depth")
+def convert_spacetodepth(node, **kwargs):
+ """Map MXNet's space_to_depth operator attributes to onnx's
+ SpaceToDepth operator and return the created node.
+ """
+ onnx = import_onnx_modules()
+ name = node["name"]
+ proc_nodes = kwargs["proc_nodes"]
+ inputs = node["inputs"]
+ attrs = node["attrs"]
+
+ input_node_id = kwargs["index_lookup"][inputs[0][0]]
+ input_node = proc_nodes[input_node_id].name
+
+ blksize = int(attrs.get("block_size", 0))
+
+ node = onnx.helper.make_node(
+ "SpaceToDepth",
+ [input_node],
+ [name],
+ blocksize=blksize,
+ name=name,
+ )
+ return [node]
diff --git a/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py
b/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py
index c44403d4992..f5a3d63f429 100644
--- a/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py
+++ b/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py
@@ -37,7 +37,7 @@
from ._op_translations import reduce_sum_square, reduce_l1, reduce_l2,
max_roi_pooling
from ._op_translations import log_softmax, softsign, lesser, greater, equal
from ._op_translations import logical_and, logical_or, logical_xor, logical_not
-from ._op_translations import mean
+from ._op_translations import mean, depthtospace, spacetodepth
# convert_map defines maps of ONNX operator names to converter
functor(callable)
# defined in the op_translations module.
@@ -140,5 +140,7 @@
'Shape' : shape,
'Gather' : gather,
'HardSigmoid' : hardsigmoid,
- 'LpPool' : lp_pooling
+ 'LpPool' : lp_pooling,
+ 'DepthToSpace' : depthtospace,
+ 'SpaceToDepth' : spacetodepth
}
diff --git a/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py
b/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py
index 5f3b756d172..7040103c005 100644
--- a/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py
+++ b/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py
@@ -688,3 +688,15 @@ def max_roi_pooling(attrs, inputs, proto_obj):
'spatial_scale':
'spatial_scale'
})
return 'ROIPooling', new_attrs, inputs
+
+def depthtospace(attrs, inputs, proto_obj):
+ """Rearranges data from depth into blocks of spatial data."""
+ new_attrs = translation_utils._fix_attribute_names(attrs,
{'blocksize':'block_size'})
+
+ return "depth_to_space", new_attrs, inputs
+
+def spacetodepth(attrs, inputs, proto_obj):
+ """Rearranges blocks of spatial data into depth."""
+ new_attrs = translation_utils._fix_attribute_names(attrs,
{'blocksize':'block_size'})
+
+ return "space_to_depth", new_attrs, inputs
diff --git a/tests/python-pytest/onnx/export/mxnet_export_test.py
b/tests/python-pytest/onnx/export/mxnet_export_test.py
index 7e1df07cbaa..7cbc9800382 100644
--- a/tests/python-pytest/onnx/export/mxnet_export_test.py
+++ b/tests/python-pytest/onnx/export/mxnet_export_test.py
@@ -25,18 +25,22 @@
from __future__ import absolute_import
import sys
import os
+import unittest
import logging
import tarfile
from collections import namedtuple
import numpy as np
import numpy.testing as npt
-from onnx import numpy_helper
+from onnx import numpy_helper, helper
from onnx import TensorProto
from mxnet.test_utils import download
from mxnet.contrib import onnx as onnx_mxnet
import mxnet as mx
CURR_PATH = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
-sys.path.insert(0, os.path.join(CURR_PATH, '../../python/unittest'))
+sys.path.insert(0, os.path.join(CURR_PATH, '../../../python/unittest'))
+import backend
+from common import with_seed
+
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
URLS = {
@@ -179,6 +183,36 @@ def test_model_accuracy(model_name, input_shape):
npt.assert_equal(expected.shape, actual.shape)
npt.assert_almost_equal(expected, actual, decimal=3)
+@with_seed()
+def test_spacetodepth():
+ n, c, h, w = shape = (1, 1, 4, 6)
+ input1 = np.random.rand(n, c, h, w).astype("float32")
+ blocksize = 2
+ inputs = [helper.make_tensor_value_info("input1", TensorProto.FLOAT,
shape=shape)]
+
+ outputs = [helper.make_tensor_value_info("output", TensorProto.FLOAT,
shape=(1, 4, 2, 3))]
+
+ nodes = [helper.make_node("SpaceToDepth", ["input1"], ["output"],
block_size=blocksize)]
+
+ graph = helper.make_graph(nodes,
+ "spacetodepth_test",
+ inputs,
+ outputs)
+
+ spacetodepth_model = helper.make_model(graph)
+
+ bkd_rep = backend.prepare(spacetodepth_model)
+ output = bkd_rep.run([input1])
+
+ tmp = np.reshape(input1, [n, c,
+ h // blocksize, blocksize,
+ w // blocksize, blocksize])
+ tmp = np.transpose(tmp, [0, 3, 5, 1, 2, 4])
+ numpy_op = np.reshape(tmp, [n, c * (blocksize**2),
+ h // blocksize,
+ w // blocksize])
+
+ npt.assert_almost_equal(output[0], numpy_op)
if __name__ == '__main__':
test_models("bvlc_googlenet", (1, 3, 224, 224), (1, 1000))
@@ -189,3 +223,5 @@ def test_model_accuracy(model_name, input_shape):
# ONNX expected results due to AveragePool issue github issue(#10194)
test_model_accuracy("inception_v1", (1, 3, 224, 224))
test_model_accuracy("inception_v2", (1, 3, 224, 224))
+
+ unittest.main()
\ No newline at end of file
diff --git a/tests/python-pytest/onnx/export/onnx_backend_test.py
b/tests/python-pytest/onnx/export/onnx_backend_test.py
index 2216a8f407f..add5e830775 100644
--- a/tests/python-pytest/onnx/export/onnx_backend_test.py
+++ b/tests/python-pytest/onnx/export/onnx_backend_test.py
@@ -91,7 +91,8 @@
'test_operator_params',
'test_operator_permute2',
'test_clip'
- 'test_cast'
+ 'test_cast',
+ 'test_depthtospace'
]
BASIC_MODEL_TESTS = [
diff --git a/tests/python-pytest/onnx/import/test_cases.py
b/tests/python-pytest/onnx/import/test_cases.py
index fd578221770..fcd01b13a83 100644
--- a/tests/python-pytest/onnx/import/test_cases.py
+++ b/tests/python-pytest/onnx/import/test_cases.py
@@ -85,7 +85,8 @@
'test_operator_exp',
'test_operator_maxpool',
'test_operator_params',
- 'test_operator_permute2'
+ 'test_operator_permute2',
+ 'test_depthtospace'
]
BASIC_MODEL_TESTS = [
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
