Merge pull request #3334 from NervanaSystems/etusien/depth_to_space

[Py] Added depth_to_space operator to Python API.

doc/sphinx/source/python_api/_autosummary/ngraph.ops.rst

@@ -31,6 +31,7 @@ ngraph.ops
       convolution_backprop_data
       cos
       cosh
+      depth_to_space
       divide
       dot
       elu

python/ngraph/__init__.py

@@ -44,6 +44,7 @@ from ngraph.ops import convolution
 from ngraph.ops import convolution_backprop_data
 from ngraph.ops import cos
 from ngraph.ops import cosh
+from ngraph.ops import depth_to_space
 from ngraph.ops import divide
 from ngraph.ops import dot
 from ngraph.ops import elu

python/ngraph/impl/op/__init__.py

@@ -68,6 +68,7 @@ from _pyngraph.op import ConvolutionBackpropData
 from _pyngraph.op import ConvolutionBackpropFilters
 from _pyngraph.op import Cos
 from _pyngraph.op import Cosh
+from _pyngraph.op import DepthToSpace
 from _pyngraph.op import Divide
 from _pyngraph.op import Dot
 from _pyngraph.op import Elu

python/ngraph/ops.py

@@ -22,11 +22,11 @@ from ngraph.impl import AxisSet, AxisVector, Coordinate, CoordinateDiff, Functio
 
 from ngraph.impl.op import Abs, Acos, Add, And, Asin, ArgMax, ArgMin, Atan, AvgPool, \
     BatchNormTraining, BatchNormInference, Broadcast, Ceiling, Clamp, Concat, Constant, Convert, \
-    Convolution, ConvolutionBackpropData, Cos, Cosh, Divide, Dot, Elu, Equal, Exp, Floor, \
-    Gelu, Gemm, GetOutputElement, Greater, GreaterEq, GRN, Less, LessEq, Log, LRN, Max, Maximum, \
-    MaxPool, Min, Minimum, Multiply, Negative, Not, NotEqual, OneHot, Or, Pad, Parameter, Product, \
-    Power, Relu, ReplaceSlice, Reshape, Reverse, Select, Sign, Sin, Sinh, Slice, Softmax, \
-    Sqrt, Subtract, Sum, Tan, Tanh, TopK
+    Convolution, ConvolutionBackpropData, Cos, Cosh, DepthToSpace, Divide, Dot, Elu, Equal, Exp, \
+    Floor, Gelu, Gemm, GetOutputElement, Greater, GreaterEq, GRN, Less, LessEq, Log, LRN, Max, \
+    Maximum, MaxPool, Min, Minimum, Multiply, Negative, Not, NotEqual, OneHot, Or, Pad, \
+    Parameter, Product, Power, Relu, ReplaceSlice, Reshape, Reverse, Select, Sign, Sin, Sinh, \
+    Slice, Softmax, Sqrt, Subtract, Sum, Tan, Tanh, TopK
 
 from typing import Callable, Iterable, List, Union
 
@@ -584,6 +584,27 @@ def convert(node, new_type, name=None):  # type: (Node, NumericType, str) -> Nod
 
 
 @nameable_op
+def depth_to_space(node, block_size, name=None):  # type: (Node, int, str) -> Node
+    """Rearranges input tensor from depth into blocks of spatial data.
+
+    Values from the height and width dimensions are moved to the depth dimension.
+
+    Input tensor has shape [N,C,H,W], where N is the batch axis, C is the channel or depth,
+    H is the height and W is the width.
+
+    Output node produces a tensor with shape:
+
+    [N, C * :code:`block_size` * :code:`block_size`, H / :code:`block_size`, W / :code:`block_size`]
+
+    :param node: The node with input tensor data.
+    :param block_size: The size of the spatial block of values describing
+                       how the tensor's data is to be rearranged.
+    :param name: Optional output node name.
+    :return: The new node performing an DepthToSpace operation on its input tensor.
+    """
+    return DepthToSpace(node, block_size)
+
+
 def gelu(node, name=None):  # type: (NodeInput, str) -> Node
     r"""Perform Gaussian Error Linear Unit operation element-wise on data from input node.
 

python/pyngraph/ops/fused/depth_to_space.cpp

+//*****************************************************************************
+// Copyright 2017-2019 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.
+//*****************************************************************************
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "ngraph/op/fused/depth_to_space.hpp"
+#include "pyngraph/ops/fused/depth_to_space.hpp"
+
+namespace py = pybind11;
+
+void regclass_pyngraph_op_DepthToSpace(py::module m)
+{
+    py::class_<ngraph::op::DepthToSpace, std::shared_ptr<ngraph::op::DepthToSpace>, ngraph::op::Op>
+        depthtospace(m, "DepthToSpace");
+    depthtospace.doc() = "ngraph.impl.op.DepthToSpace wraps ngraph::op::DepthToSpace";
+    depthtospace.def(py::init<const std::shared_ptr<ngraph::Node>&, std::size_t&>());
+}

python/pyngraph/ops/fused/depth_to_space.hpp

+//*****************************************************************************
+// Copyright 2017-2019 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.
+//*****************************************************************************
+
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+
+void regclass_pyngraph_op_DepthToSpace(py::module m);

python/pyngraph/ops/regmodule_pyngraph_op.cpp

@@ -48,6 +48,7 @@ void regmodule_pyngraph_op(py::module m_op)
     regclass_pyngraph_op_ConvolutionBackpropFilters(m_op);
     regclass_pyngraph_op_Cos(m_op);
     regclass_pyngraph_op_Cosh(m_op);
+    regclass_pyngraph_op_DepthToSpace(m_op);
     regclass_pyngraph_op_Divide(m_op);
     regclass_pyngraph_op_Dot(m_op);
     regclass_pyngraph_op_Elu(m_op);

python/pyngraph/ops/regmodule_pyngraph_op.hpp

@@ -43,6 +43,7 @@
 #include "pyngraph/ops/exp.hpp"
 #include "pyngraph/ops/floor.hpp"
 #include "pyngraph/ops/fused/clamp.hpp"
+#include "pyngraph/ops/fused/depth_to_space.hpp"
 #include "pyngraph/ops/fused/elu.hpp"
 #include "pyngraph/ops/fused/gelu.hpp"
 #include "pyngraph/ops/fused/gemm.hpp"

python/setup.py

@@ -178,6 +178,7 @@ sources = [
     'pyngraph/ops/cos.cpp',
     'pyngraph/ops/cosh.cpp',
     'pyngraph/ops/ceiling.cpp',
+    'pyngraph/ops/fused/depth_to_space.cpp',
     'pyngraph/ops/divide.cpp',
     'pyngraph/ops/dot.cpp',
     'pyngraph/ops/fused/elu.cpp',

python/test/ngraph/test_ops_fused.py

@@ -69,6 +69,33 @@ def test_elu_operator_with_scalar():
     assert np.allclose(result, expected)
 
 
+def test_depth_to_space():
+    runtime = get_runtime()
+
+    data_value = np.array([[[[0, 1, 2],
+                            [3, 4, 5]],
+                            [[6, 7, 8],
+                            [9, 10, 11]],
+                            [[12, 13, 14],
+                            [15, 16, 17]],
+                            [[18, 19, 20],
+                            [21, 22, 23]]]], dtype=np.float32)
+    block_size = np.float32(2)
+
+    data_shape = [1, 4, 2, 3]
+    parameter_data = ng.parameter(data_shape, name='Data', dtype=np.float32)
+
+    model = ng.depth_to_space(parameter_data, block_size)
+    computation = runtime.computation(model, parameter_data)
+
+    result = computation(data_value)
+    expected = np.array([[[[0, 6, 1, 7, 2, 8],
+                        [12, 18, 13, 19, 14, 20],
+                        [3, 9, 4, 10, 5, 11],
+                        [15, 21, 16, 22, 17, 23]]]], dtype=np.float32)
+    assert np.allclose(result, expected)
+
+
 def test_gemm_operator():
     runtime = get_runtime()