This is an automated email from the ASF dual-hosted git repository.
haibin pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push:
new c13ce5c [MXNET-378] Adding depth_to_space and space_to_depth
operator(Updated) (#11587)
c13ce5c is described below
commit c13ce5cb0adbef4359ced3c07f804f7172ab4c8d
Author: access2rohit <[email protected]>
AuthorDate: Thu Jul 26 13:44:23 2018 -0700
[MXNET-378] Adding depth_to_space and space_to_depth operator(Updated)
(#11587)
* [MXNET-378] Adding depth_to_space and space_to_depth operator
* fixed lint and windows CPU errors
* compliance with C++ style guiide and address shortcomings in unittests
* fixed documentation and nitpicky suggestions
* added operator references in API docs and removed inplace optimization
support
* Added references in symbol.md and ndarray.md. Improved test cases and
added block_size check
* Fixing bugs in documentation. Tests now include tensors of random shapes.
---
docs/api/python/ndarray/ndarray.md | 4 +
docs/api/python/symbol/symbol.md | 4 +
python/mxnet/ndarray/ndarray.py | 16 ++
python/mxnet/symbol/symbol.py | 16 ++
src/operator/tensor/matrix_op-inl.h | 322 +++++++++++++++++++++++++++++++++
src/operator/tensor/matrix_op.cc | 107 +++++++++++
src/operator/tensor/matrix_op.cu | 6 +
tests/python/unittest/test_operator.py | 100 ++++++++++
8 files changed, 575 insertions(+)
diff --git a/docs/api/python/ndarray/ndarray.md
b/docs/api/python/ndarray/ndarray.md
index d92c3e8..01a1544 100644
--- a/docs/api/python/ndarray/ndarray.md
+++ b/docs/api/python/ndarray/ndarray.md
@@ -156,6 +156,8 @@ The `ndarray` package provides several classes:
NDArray.transpose
NDArray.swapaxes
NDArray.flip
+ NDArray.depth_to_space
+ NDArray.space_to_depth
```
### Array reduction
@@ -411,6 +413,8 @@ The `ndarray` package provides several classes:
transpose
swapaxes
flip
+ depth_to_space
+ space_to_depth
```
### Joining and splitting arrays
diff --git a/docs/api/python/symbol/symbol.md b/docs/api/python/symbol/symbol.md
index b0db774..7c78cbd 100644
--- a/docs/api/python/symbol/symbol.md
+++ b/docs/api/python/symbol/symbol.md
@@ -222,6 +222,8 @@ Composite multiple symbols into a new one by an operator.
Symbol.transpose
Symbol.swapaxes
Symbol.flip
+ Symbol.depth_to_space
+ Symbol.space_to_depth
```
### Reduce functions
@@ -409,6 +411,8 @@ Composite multiple symbols into a new one by an operator.
transpose
swapaxes
flip
+ depth_to_space
+ space_to_depth
```
### Joining and splitting symbols
diff --git a/python/mxnet/ndarray/ndarray.py b/python/mxnet/ndarray/ndarray.py
index 64d5102..46b21a9 100644
--- a/python/mxnet/ndarray/ndarray.py
+++ b/python/mxnet/ndarray/ndarray.py
@@ -1302,6 +1302,22 @@ fixed-size items.
"""
return op.flip(self, *args, **kwargs)
+ def depth_to_space(self, *args, **kwargs):
+ """Convenience fluent method for :py:func:`depth_to_space`.
+
+ The arguments are the same as for :py:func:`depth_to_space`, with
+ this array as data.
+ """
+ return op.depth_to_space(self, *args, **kwargs)
+
+ def space_to_depth(self, *args, **kwargs):
+ """Convenience fluent method for :py:func:`space_to_depth`.
+
+ The arguments are the same as for :py:func:`space_to_depth`, with
+ this array as data.
+ """
+ return op.space_to_depth(self, *args, **kwargs)
+
def diag(self, k=0, **kwargs):
"""Convenience fluent method for :py:func:`diag`.
diff --git a/python/mxnet/symbol/symbol.py b/python/mxnet/symbol/symbol.py
index ea476cd..5f6cbd6 100644
--- a/python/mxnet/symbol/symbol.py
+++ b/python/mxnet/symbol/symbol.py
@@ -2046,6 +2046,22 @@ class Symbol(SymbolBase):
"""
return op.flip(self, *args, **kwargs)
+ def depth_to_space(self, *args, **kwargs):
+ """Convenience fluent method for :py:func:`depth_to_space`.
+
+ The arguments are the same as for :py:func:`depth_to_space`, with
+ this array as data.
+ """
+ return op.depth_to_space(self, *args, **kwargs)
+
+ def space_to_depth(self, *args, **kwargs):
+ """Convenience fluent method for :py:func:`space_to_depth`.
+
+ The arguments are the same as for :py:func:`space_to_depth`, with
+ this array as data.
+ """
+ return op.space_to_depth(self, *args, **kwargs)
+
def diag(self, k=0, **kwargs):
"""Convenience fluent method for :py:func:`diag`.
diff --git a/src/operator/tensor/matrix_op-inl.h
b/src/operator/tensor/matrix_op-inl.h
index dcdf03a..eec9205 100644
--- a/src/operator/tensor/matrix_op-inl.h
+++ b/src/operator/tensor/matrix_op-inl.h
@@ -2171,6 +2171,328 @@ inline bool SqueezeShape(const nnvm::NodeAttrs& attrs,
return true;
}
+struct DepthToSpaceParam : public dmlc::Parameter<DepthToSpaceParam> {
+ int block_size;
+ DMLC_DECLARE_PARAMETER(DepthToSpaceParam) {
+ DMLC_DECLARE_FIELD(block_size)
+ .describe("Blocks of [block_size. block_size] are moved");
+ }
+};
+
+inline bool DepthToSpaceOpShape(const nnvm::NodeAttrs& attrs,
+ std::vector<TShape>* in_attrs,
+ std::vector<TShape>* out_attrs) {
+ const DepthToSpaceParam& param = nnvm::get<DepthToSpaceParam>(attrs.parsed);
+ CHECK_EQ(in_attrs->size(), 1U);
+ CHECK_EQ(out_attrs->size(), 1U);
+ CHECK_EQ(in_attrs->at(0).ndim(), 4) << "Operation Depth To Space requires
exactly 4D tensor";
+
+ TShape expected_out(4);
+
+ TShape& in_shape = in_attrs->at(0);
+ int block = param.block_size;
+ CHECK_NE(block, 0) << "block_size must be a positive integer value";
+ CHECK_NE(in_shape[1], 0) << "Depth dimension:1 cannot be 0";
+ CHECK_EQ(in_shape[1] % (block * block), 0)
+ << "Cannot perform Depth To Space operation on the specified tensor."
+ " Dimension:1(depth dimension) should be a multiple of 'block^2'";
+ CHECK_NE(in_shape[0], 0)
+ << "Operation requires a 4D tensor. Size of dimension:0 cannot be 0";
+ CHECK_NE(in_shape[2], 0)
+ << "Operation requires a 4D tensor. Size of dimension:2 cannot be 0";
+ CHECK_NE(in_shape[3], 0)
+ << "Operation requires a 4D tensor. Size of dimension:3 cannot be 0";
+
+ expected_out[0] = in_shape[0];
+ expected_out[1] = in_shape[1] / (block * block);
+ uint32_t i = 2;
+ while (i < expected_out.ndim()) {
+ expected_out[i] = in_shape[i] * block;
+ ++i;
+ }
+
+ SHAPE_ASSIGN_CHECK(*out_attrs, 0, expected_out);
+ return true;
+}
+
+inline bool DepthToSpaceOpType(const nnvm::NodeAttrs& attrs,
+ std::vector<int>* in_attrs,
+ std::vector<int>* out_attrs) {
+ CHECK_EQ(in_attrs->size(), 1U);
+ CHECK_EQ(out_attrs->size(), 1U);
+
+ TYPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(0));
+ TYPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(0));
+ return out_attrs->at(0) != -1;
+}
+
+/*!
+ * \brief This function updates the value of input index from where the data
element
+ * needs to be fetched and written out to the ith location in output tensor
+ * \param index_position index within offset array to get offset of given
dimension
+ * \param dim_size size of current dimension
+ * \param idx output tensor index
+ * \param inp_index index within input tensor from where value is
retrieved
+ * \param offset_arr array containing the linear offset of input tensor
+ */
+MSHADOW_XINLINE void update_index(int index_position, int dim_size, int *idx,
+ int *inp_index, const int* offset_arr) {
+ int next_idx_val = *idx / dim_size;
+ *inp_index += (*idx - next_idx_val * dim_size) * offset_arr[index_position];
+ *idx = next_idx_val;
+}
+
+/*!
+ * \brief This function performs the tensor transpose (0, 1, 2, 3, 4, 5) ->
+ * (0, 3, 4, 1, 5, 2) by computing linear index within input tensor to be
mapped
+ * to the ith index of output tensor
+ * \param i tensor index
+ * \param out_data output tensor
+ * \param in_data input tensor
+ * \param block size of chunks to be moved out of depth dimension
+ * \param size array containing the size of each dimension of input
tensor
+ * \param offset_arr array containing the linear offset of input tensor
+ */
+template<int req>
+struct depth_to_space_forward {
+ template<typename DType>
+ MSHADOW_XINLINE static void Map(int i, DType* out_data, const DType* in_data,
+ const int block, const int* size, const int*
offset_arr) {
+ int inp_index = 0, idx = i, dim_size;
+ dim_size = block;
+ update_index(2, dim_size, &idx, &inp_index, offset_arr);
+ dim_size = size[3];
+ update_index(5, dim_size, &idx, &inp_index, offset_arr);
+ dim_size = block;
+ update_index(1, dim_size, &idx, &inp_index, offset_arr);
+ dim_size = size[2];
+ update_index(4, dim_size, &idx, &inp_index, offset_arr);
+ dim_size = size[1] / (block * block);
+ update_index(3, dim_size, &idx, &inp_index, offset_arr);
+ dim_size = size[0];
+ update_index(0, dim_size, &idx, &inp_index, offset_arr);
+ KERNEL_ASSIGN(out_data[i], req, in_data[inp_index]);
+ }
+};
+
+/*!
+ * \brief This function calculates the linear offset for each dimension of
+ * input tensor and stores them in an array, which is later used in
+ * performing depth_to_space operation
+ * \param i global thread id
+ * \param offset_arr array to be populated with offset values
+ * \param size array to be populated with size of each dimension of
input tensor
+ * \param block size of chunks to be moved out of depth dimension
+ * \param size0 size of Dim 0 of input tensor
+ * \param size1 size of Dim 1 of input tensor
+ * \param size2 size of Dim 2 of input tensor
+ * \param size3 size of Dim 3 of input tensor
+ */
+template<int req>
+struct compute_offset_for_depth_to_space {
+ template<typename DType>
+ MSHADOW_XINLINE static void Map(int i, DType* offset_arr, DType* size, const
int block,
+ const int32_t size0, const int32_t size1,
const int32_t size2,
+ const int32_t size3) {
+ size[0] = size0;
+ size[1] = size1;
+ size[2] = size2;
+ size[3] = size3;
+
+ offset_arr[5] = 1;
+ offset_arr[4] = offset_arr[5] * size[3];
+ offset_arr[3] = offset_arr[4] * size[2];
+ offset_arr[2] = offset_arr[3] * size[1] / (block * block);
+ offset_arr[1] = offset_arr[2] * block;
+ offset_arr[0] = offset_arr[1] * block;
+ }
+};
+
+template<typename xpu>
+void DepthToSpaceOpForward(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<TBlob>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<TBlob>& outputs) {
+ CHECK_EQ(inputs.size(), 1U);
+ CHECK_EQ(outputs.size(), 1U);
+ CHECK_EQ(req.size(), 1U);
+ mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
+ const TBlob& in_data = inputs[0];
+ const TBlob& out_data = outputs[0];
+ const DepthToSpaceParam& param = nnvm::get<DepthToSpaceParam>(attrs.parsed);
+ using namespace mxnet_op;
+ int block = param.block_size;
+
+ mshadow::Tensor<xpu, 1, char> workspace =
+ ctx.requested[0].get_space_typed<xpu, 1,
char>(mshadow::Shape1(sizeof(int32_t) * 10), s);
+ char* workspace_curr_ptr = workspace.dptr_;
+ int32_t* offset_arr = reinterpret_cast<int32_t*>(workspace_curr_ptr);
+ int32_t* size = reinterpret_cast<int32_t*>(workspace_curr_ptr +
sizeof(int32_t) * 6);
+
+ MSHADOW_TYPE_SWITCH(out_data.type_flag_, DType, {
+ MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
+ Kernel<compute_offset_for_depth_to_space<req_type>, xpu>::Launch(
+ s, 1, offset_arr, size, block, in_data.shape_[0], in_data.shape_[1],
+ in_data.shape_[2], in_data.shape_[3]);
+
+ Kernel<depth_to_space_forward<req_type>, xpu>::Launch(
+ s, out_data.Size(), out_data.dptr<DType>(), in_data.dptr<DType>(),
+ block, size, offset_arr);
+ });
+ });
+}
+
+inline bool SpaceToDepthOpShape(const nnvm::NodeAttrs& attrs,
+ std::vector<TShape>* in_attrs,
+ std::vector<TShape>* out_attrs) {
+ const DepthToSpaceParam& param = nnvm::get<DepthToSpaceParam>(attrs.parsed);
+ CHECK_EQ(in_attrs->size(), 1U);
+ CHECK_EQ(out_attrs->size(), 1U);
+ CHECK_EQ(in_attrs->at(0).ndim(), 4) << "Operation Space To Depth requires
exactly 4D tensor";
+
+ TShape expected_out(in_attrs->at(0).ndim());
+
+ TShape& in_shape = in_attrs->at(0);
+ int block = param.block_size;
+ CHECK_NE(block, 0) << "block_size must be a positive integer value";
+ CHECK_NE(in_shape[0], 0)
+ << "Operation requires a 4D tensor. Size of dimension:0 cannot be 0";
+ CHECK_NE(in_shape[1], 0) << "Depth dimension:1 cannot be 0";
+ CHECK_NE(in_shape[2], 0)
+ << "Operation requires a 4D tensor. Size of dimension:2 cannot be 0";
+ CHECK_EQ(in_shape[2] % block, 0)
+ << "Cannot perform Depth To Space operation on the specified tensor."
+ " Dimension:2(1st Space dimension) should be a multiple of 'block' ";
+ CHECK_NE(in_shape[3], 0)
+ << "Operation requires a 4D tensor. Size of dimension:3 cannot be 0";
+ CHECK_EQ(in_shape[3] % block, 0)
+ << "Cannot perform Depth To Space operation on the specified tensor."
+ " Dimension:3(2nd space dimension) should be a multiple of 'block' ";
+
+ expected_out[0] = in_shape[0];
+ expected_out[1] = in_shape[1] * block * block;
+ uint32_t i = 2;
+ while (i < expected_out.ndim()) {
+ expected_out[i] = in_shape[i] / block;
+ ++i;
+ }
+
+ SHAPE_ASSIGN_CHECK(*out_attrs, 0, expected_out);
+ return true;
+}
+
+inline bool SpaceToDepthOpType(const nnvm::NodeAttrs& attrs,
+ std::vector<int>* in_attrs,
+ std::vector<int>* out_attrs) {
+ CHECK_EQ(in_attrs->size(), 1U);
+ CHECK_EQ(out_attrs->size(), 1U);
+
+ TYPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(0));
+ TYPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(0));
+ return out_attrs->at(0) != -1;
+}
+
+/*!
+ * \brief This function preforms the tensor transpose (0, 1, 2, 3, 4, 5) ->
+ * (0, 3, 5, 1, 2, 4) by computing linear index within input tensor to be
mapped
+ * to the ith index of output tensor
+ * \param i tensor index
+ * \param out_data output tensor
+ * \param in_data input tensor
+ * \param block size of chunks to be moved out of depth dimension
+ * \param size array containing the size of each dimension of input
tensor
+ * \param offset_arr array containing the linear offset of input tensor
+ */
+template<int req>
+struct space_to_depth_forward {
+ template<typename DType>
+ MSHADOW_XINLINE static void Map(int i, DType* out_data, const DType*
in_data, const int block,
+ const int* size, const int* offset_arr) {
+ int inp_index = 0, idx = i, dim_size;
+ dim_size = size[3] / block;
+ update_index(4, dim_size, &idx, &inp_index, offset_arr);
+ dim_size = size[2] / block;
+ update_index(2, dim_size, &idx, &inp_index, offset_arr);
+ dim_size = size[1];
+ update_index(1, dim_size, &idx, &inp_index, offset_arr);
+ dim_size = block;
+ update_index(5, dim_size, &idx, &inp_index, offset_arr);
+ dim_size = block;
+ update_index(3, dim_size, &idx, &inp_index, offset_arr);
+ dim_size = size[0];
+ update_index(0, dim_size, &idx, &inp_index, offset_arr);
+ KERNEL_ASSIGN(out_data[i], req, in_data[inp_index]);
+ }
+};
+
+/*!
+ * \brief This function calculates the linear offset for each dimension of
+ * input tensor and stores them in an array, which is later used in
+ * performing space_to_depth operation
+ * \param i global thread id
+ * \param offset_arr array to be populated with offset values
+ * \param size array to be populated with size of each dimension of
input tensor
+ * \param block size of chunks to be moved out of depth dimension
+ * \param size0 size of Dim 0 of input tensor
+ * \param size1 size of Dim 1 of input tensor
+ * \param size2 size of Dim 2 of input tensor
+ * \param size3 size of Dim 3 of input tensor
+ */
+template<int req>
+struct compute_offset_for_space_to_depth {
+ template<typename DType>
+ MSHADOW_XINLINE static void Map(int i, DType* offset_arr, DType* size, const
int block,
+ const int32_t size0, const int32_t size1,
+ const int32_t size2, const int32_t size3) {
+ size[0] = size0;
+ size[1] = size1;
+ size[2] = size2;
+ size[3] = size3;
+
+ offset_arr[5] = 1;
+ offset_arr[4] = offset_arr[5] * block;
+ offset_arr[3] = offset_arr[4] * size[3] / block;
+ offset_arr[2] = offset_arr[3] * block;
+ offset_arr[1] = offset_arr[2] * size[2] / block;
+ offset_arr[0] = offset_arr[1] * size[1];
+ }
+};
+
+template<typename xpu>
+void SpaceToDepthOpForward(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<TBlob>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<TBlob>& outputs) {
+ CHECK_EQ(inputs.size(), 1U);
+ CHECK_EQ(outputs.size(), 1U);
+ CHECK_EQ(req.size(), 1U);
+ mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
+ const TBlob& in_data = inputs[0];
+ const TBlob& out_data = outputs[0];
+ const DepthToSpaceParam& param = nnvm::get<DepthToSpaceParam>(attrs.parsed);
+ using namespace mxnet_op;
+ int block = param.block_size;
+
+ mshadow::Tensor<xpu, 1, char> workspace =
+ ctx.requested[0].get_space_typed<xpu, 1,
char>(mshadow::Shape1(sizeof(int32_t) * 10), s);
+ char* workspace_curr_ptr = workspace.dptr_;
+ int32_t* offset_arr = reinterpret_cast<int32_t*>(workspace_curr_ptr);
+ int32_t* size = reinterpret_cast<int32_t*>(workspace_curr_ptr +
sizeof(int32_t) * 6);
+
+ MSHADOW_TYPE_SWITCH(out_data.type_flag_, DType, {
+ MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
+ Kernel<compute_offset_for_space_to_depth<req_type>, xpu>::Launch(
+ s, 1, offset_arr, size, block, in_data.shape_[0], in_data.shape_[1],
+ in_data.shape_[2], in_data.shape_[3]);
+ Kernel<space_to_depth_forward<req_type>, xpu>::Launch(
+ s, out_data.Size(), out_data.dptr<DType>(), in_data.dptr<DType>(),
+ block, size, offset_arr);
+ });
+ });
+}
+
} // namespace op
} // namespace mxnet
diff --git a/src/operator/tensor/matrix_op.cc b/src/operator/tensor/matrix_op.cc
index 29d493a..ffdc228 100644
--- a/src/operator/tensor/matrix_op.cc
+++ b/src/operator/tensor/matrix_op.cc
@@ -101,6 +101,7 @@ DMLC_REGISTER_PARAMETER(TileParam);
DMLC_REGISTER_PARAMETER(ReverseParam);
DMLC_REGISTER_PARAMETER(StackParam);
DMLC_REGISTER_PARAMETER(SqueezeParam);
+DMLC_REGISTER_PARAMETER(DepthToSpaceParam);
NNVM_REGISTER_OP(Reshape)
.add_alias("reshape")
@@ -908,5 +909,111 @@ NNVM_REGISTER_OP(_backward_squeeze)
.set_attr<nnvm::TIsBackward>("TIsBackward", true)
.set_attr<FCompute>("FCompute<cpu>", UnaryOp::IdentityCompute<cpu>);
+NNVM_REGISTER_OP(depth_to_space)
+.describe(R"code(Rearranges(permutes) data from depth into blocks of spatial
data.
+Similar to ONNX DepthToSpace operator:
+https://github.com/onnx/onnx/blob/master/docs/Operators.md#DepthToSpace.
+The output is a new tensor where the values from depth dimension are moved in
spatial blocks
+to height and width dimension. The reverse of this operation is
``space_to_depth``.
+
+.. math::
+
+ \begin{gather*}
+ x \prime = reshape(x, [N, block\_size, block\_size, C / (block\_size ^ 2),
H * block\_size, W * block\_size]) \\
+ x \prime \prime = transpose(x \prime, [0, 3, 4, 1, 5, 2]) \\
+ y = reshape(x \prime \prime, [N, C / (block\_size ^ 2), H * block\_size, W
* block\_size])
+ \end{gather*}
+
+where :math:`x` is an input tensor with default layout as :math:`[N, C, H,
W]`: [batch, channels, height, width]
+and :math:`y` is the output tensor of layout :math:`[N, C / (block\_size ^ 2),
H * block\_size, W * block\_size]`
+
+Example::
+
+ x = [[[[0, 1, 2],
+ [3, 4, 5]],
+ [[6, 7, 8],
+ [9, 10, 11]],
+ [[12, 13, 14],
+ [15, 16, 17]],
+ [[18, 19, 20],
+ [21, 22, 23]]]]
+
+ depth_to_space(x, 2) = [[[[0, 6, 1, 7, 2, 8],
+ [12, 18, 13, 19, 14, 20],
+ [3, 9, 4, 10, 5, 11],
+ [15, 21, 16, 22, 17, 23]]]]
+)code" ADD_FILELINE)
+.set_attr_parser(ParamParser<DepthToSpaceParam>)
+.set_num_inputs(1)
+.set_num_outputs(1)
+.set_attr<nnvm::FListInputNames>("FListInputNames",
+ [](const NodeAttrs& attrs) {
+ return std::vector<std::string>{"data"};
+ })
+.set_attr<nnvm::FInferShape>("FInferShape", DepthToSpaceOpShape)
+.set_attr<nnvm::FInferType>("FInferType", DepthToSpaceOpType)
+.set_attr<FCompute>("FCompute<cpu>", DepthToSpaceOpForward<cpu>)
+.set_attr<FResourceRequest>("FResourceRequest",
+ [](const NodeAttrs& n) {
+ return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
+})
+.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"space_to_depth"})
+.add_argument("data", "NDArray-or-Symbol", "Input ndarray")
+.add_arguments(DepthToSpaceParam::__FIELDS__());
+
+NNVM_REGISTER_OP(space_to_depth)
+.describe(R"code(Rearranges(permutes) blocks of spatial data into depth.
+Similar to ONNX SpaceToDepth operator:
+https://github.com/onnx/onnx/blob/master/docs/Operators.md#SpaceToDepth
+
+The output is a new tensor where the values from height and width dimension
are
+moved to the depth dimension. The reverse of this operation is
``depth_to_space``.
+
+.. math::
+
+ \begin{gather*}
+ x \prime = reshape(x, [N, C, H / block\_size, block\_size, W /
block\_size, block\_size]) \\
+ x \prime \prime = transpose(x \prime, [0, 3, 5, 1, 2, 4]) \\
+ y = reshape(x \prime \prime, [N, C * (block\_size ^ 2), H / block\_size, W
/ block\_size])
+ \end{gather*}
+
+where :math:`x` is an input tensor with default layout as :math:`[N, C, H,
W]`: [batch, channels, height, width]
+and :math:`y` is the output tensor of layout :math:`[N, C * (block\_size ^ 2),
H / block\_size, W / block\_size]`
+
+Example::
+
+ x = [[[[0, 6, 1, 7, 2, 8],
+ [12, 18, 13, 19, 14, 20],
+ [3, 9, 4, 10, 5, 11],
+ [15, 21, 16, 22, 17, 23]]]]
+
+
+ space_to_depth(x, 2) = [[[[0, 1, 2],
+ [3, 4, 5]],
+ [[6, 7, 8],
+ [9, 10, 11]],
+ [[12, 13, 14],
+ [15, 16, 17]],
+ [[18, 19, 20],
+ [21, 22, 23]]]]
+)code" ADD_FILELINE)
+.set_attr_parser(ParamParser<DepthToSpaceParam>)
+.set_num_inputs(1)
+.set_num_outputs(1)
+.set_attr<nnvm::FListInputNames>("FListInputNames",
+ [](const NodeAttrs& attrs) {
+ return std::vector<std::string>{"data"};
+ })
+.set_attr<nnvm::FInferShape>("FInferShape", SpaceToDepthOpShape)
+.set_attr<nnvm::FInferType>("FInferType", SpaceToDepthOpType)
+.set_attr<FCompute>("FCompute<cpu>", SpaceToDepthOpForward<cpu>)
+.set_attr<FResourceRequest>("FResourceRequest",
+ [](const NodeAttrs& n) {
+ return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
+})
+.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"depth_to_space"})
+.add_argument("data", "NDArray-or-Symbol", "Input ndarray")
+.add_arguments(DepthToSpaceParam::__FIELDS__());
+
} // namespace op
} // namespace mxnet
diff --git a/src/operator/tensor/matrix_op.cu b/src/operator/tensor/matrix_op.cu
index bd1b9f2..4e31a4c 100644
--- a/src/operator/tensor/matrix_op.cu
+++ b/src/operator/tensor/matrix_op.cu
@@ -211,5 +211,11 @@ NNVM_REGISTER_OP(squeeze)
NNVM_REGISTER_OP(_backward_squeeze)
.set_attr<FCompute>("FCompute<gpu>", UnaryOp::IdentityCompute<gpu>);
+NNVM_REGISTER_OP(depth_to_space)
+.set_attr<FCompute>("FCompute<gpu>", DepthToSpaceOpForward<gpu>);
+
+NNVM_REGISTER_OP(space_to_depth)
+.set_attr<FCompute>("FCompute<gpu>", SpaceToDepthOpForward<gpu>);
+
} // namespace op
} // namespace mxnet
diff --git a/tests/python/unittest/test_operator.py
b/tests/python/unittest/test_operator.py
index e50f8a1..fa5de0c 100644
--- a/tests/python/unittest/test_operator.py
+++ b/tests/python/unittest/test_operator.py
@@ -6679,6 +6679,106 @@ def test_diag():
diag_sym = mx.sym.diag(data=data, k=-1)
check_numeric_gradient(diag_sym, [a_np])
+@with_seed()
+def test_depthtospace():
+ def f(x, blocksize):
+ b, c, h, w = x.shape[0], x.shape[1], x.shape[2], x.shape[3]
+ tmp = np.reshape(x, [b, blocksize, blocksize, c // (blocksize**2), h,
w])
+ tmp = np.transpose(tmp, [0, 3, 4, 1, 5, 2])
+ y = np.reshape(tmp, [b, c // (blocksize**2), h * blocksize, w *
blocksize])
+ return y
+
+ block = random.randint(2, 4)
+ rand_mul1 = random.randint(1, 4)
+ n = random.randint(1, 5)
+ c = block * block * rand_mul1
+ h = random.randint(1, 5)
+ w = random.randint(1, 5)
+ shape_inp = (n, c, h, w)
+ data = rand_ndarray(shape_inp, 'default')
+ data_np = data.asnumpy()
+ expected = f(data_np, block)
+ output = mx.nd.depth_to_space(data, block)
+ assert_almost_equal(output.asnumpy(), expected, atol=1e-3, rtol=1e-3)
+
+ shape_out = (n, c // (block ** 2), h * block, w * block)
+ data = mx.sym.Variable('data')
+ dts_sym = mx.sym.depth_to_space(data, block)
+ check_numeric_gradient(dts_sym, [np.ones(shape_inp)])
+
+ check_symbolic_forward(dts_sym, [data_np], [expected])
+ check_symbolic_backward(dts_sym, [data_np], [np.ones(shape_out)],
[np.ones(shape_inp)])
+
+ def test_invalid_depth_dim():
+ invalid_shape_inp = (n, block - 1, h, w)
+ data = rand_ndarray(invalid_shape_inp, 'default')
+ assertRaises(MXNetError, mx.nd.depth_to_space, data, block)
+
+ def test_invalid_space_dim():
+ invalid_shape_inp = (n, block ** 2, 0, block + 1)
+ data = rand_ndarray(invalid_shape_inp, 'default')
+ assertRaises(MXNetError, mx.nd.depth_to_space, data, block)
+
+ def test_invalid_block_size():
+ block = 0
+ invalid_shape_inp = (n , c, h, w)
+ data = rand_ndarray(invalid_shape_inp, 'default')
+ assertRaises(MXNetError, mx.nd.depth_to_space, data, block)
+
+ test_invalid_depth_dim()
+ test_invalid_space_dim()
+ test_invalid_block_size()
+
+@with_seed()
+def test_spacetodepth():
+ def f(x, blocksize):
+ b, c, h, w = x.shape[0], x.shape[1], x.shape[2], x.shape[3]
+ tmp = np.reshape(x, [b, c, h // blocksize, blocksize, w // blocksize,
blocksize])
+ tmp = np.transpose(tmp, [0, 3, 5, 1, 2, 4])
+ y = np.reshape(tmp, [b, c * (blocksize**2), h // blocksize, w //
blocksize])
+ return y
+
+ block = random.randint(2, 4)
+ rand_mul1 = random.randint(1, 4)
+ rand_mul2 = random.randint(1, 4)
+ n = random.randint(1, 5)
+ c = random.randint(1, 5)
+ h = block * rand_mul1
+ w = block * rand_mul2
+ shape_inp = (n, c, h, w)
+ data = rand_ndarray(shape_inp, 'default')
+ data_np = data.asnumpy()
+ expected = f(data_np, block)
+ output = mx.nd.space_to_depth(data, block)
+ assert_almost_equal(output.asnumpy(), expected, atol=1e-3, rtol=1e-3)
+
+ shape_out = (n, c * (block ** 2), h // block, w // block)
+ data = mx.sym.Variable('data')
+ dts_sym = mx.sym.space_to_depth(data, block)
+ check_numeric_gradient(dts_sym, [np.ones(shape_inp)])
+
+ check_symbolic_forward(dts_sym, [data_np], [expected])
+ check_symbolic_backward(dts_sym, [data_np], [np.ones(shape_out)],
[np.ones(shape_inp)])
+
+ def test_invalid_space_dim():
+ invalid_shape_inp = (n , c, block - 1, w)
+ data = rand_ndarray(invalid_shape_inp, 'default')
+ assertRaises(MXNetError, mx.nd.space_to_depth, data, block)
+
+ def test_invalid_block_size():
+ block = 0
+ invalid_shape_inp = (n, c, h, w)
+ data = rand_ndarray(invalid_shape_inp, 'default')
+ assertRaises(MXNetError, mx.nd.space_to_depth, data, block)
+
+ def test_invalid_depth_dim():
+ invalid_shape_inp = (n, 0, h, w)
+ data = rand_ndarray(invalid_shape_inp, 'default')
+ assertRaises(MXNetError, mx.nd.space_to_depth, data, block)
+
+ test_invalid_space_dim()
+ test_invalid_block_size()
+ test_invalid_depth_dim()
if __name__ == '__main__':
import nose
