eric-haibin-lin closed pull request #11587: [MXNET-378] Adding depth_to_space
and space_to_depth operator(Updated)
URL: https://github.com/apache/incubator-mxnet/pull/11587
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/docs/api/python/ndarray/ndarray.md
b/docs/api/python/ndarray/ndarray.md
index d92c3e84e00..01a154405cf 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 b0db774d933..7c78cbd59b0 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 64d510296ff..46b21a90d4c 100644
--- a/python/mxnet/ndarray/ndarray.py
+++ b/python/mxnet/ndarray/ndarray.py
@@ -1302,6 +1302,22 @@ def flip(self, *args, **kwargs):
"""
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 ea476cdcb21..5f6cbd6b6e1 100644
--- a/python/mxnet/symbol/symbol.py
+++ b/python/mxnet/symbol/symbol.py
@@ -2046,6 +2046,22 @@ def flip(self, *args, **kwargs):
"""
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 dcdf03a5316..eec920555ed 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 29d493ae5a5..ffdc228b2d6 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 bd1b9f20825..4e31a4cf115 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 11180ebbc5d..6b0e588889a 100644
--- a/tests/python/unittest/test_operator.py
+++ b/tests/python/unittest/test_operator.py
@@ -6678,6 +6678,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
----------------------------------------------------------------
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
