This is an automated email from the ASF dual-hosted git repository.
zhasheng 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 0c0457a Add a contrib operator for dynamic reshape (#15872)
0c0457a is described below
commit 0c0457a8f7171b24aa2b9a534c0d0e3b15207147
Author: Vandana Kannan <[email protected]>
AuthorDate: Tue Sep 8 00:06:55 2020 -0700
Add a contrib operator for dynamic reshape (#15872)
* Reshape which accepts shape as a tensor
* Unit tests for dynamic reshape
* Get output shape from shape tensor
* Make changes to backward
* fix lint errors
* add newline
* handle negative values for shape
* edit args_grad shape
* Using out in forward
* Add an example
* Add a comment in forward pass
* address review comments
* add some more doc and ndarray tests
* Re-trigger CI
* fix type params
* Update test_contrib_operator.py
* Update dynamic_shape_ops.cc
* Update dynamic_shape_ops.cc
* add req support
Co-authored-by: Sheng Zha <[email protected]>
---
src/operator/contrib/dynamic_shape_ops-inl.h | 92 +++++++++++++++++
src/operator/contrib/dynamic_shape_ops.cc | 138 +++++++++++++++++++++++++
tests/python/unittest/test_contrib_operator.py | 50 +++++++++
3 files changed, 280 insertions(+)
diff --git a/src/operator/contrib/dynamic_shape_ops-inl.h
b/src/operator/contrib/dynamic_shape_ops-inl.h
new file mode 100644
index 0000000..1d1aff8
--- /dev/null
+++ b/src/operator/contrib/dynamic_shape_ops-inl.h
@@ -0,0 +1,92 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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.
+ */
+/*!
+ * Copyright (c) 2018 by Contributors
+ * \file dynamic_shape_ops-inl.h
+*/
+
+#ifndef MXNET_OPERATOR_CONTRIB_DYNAMIC_SHAPE_OPS_INL_H_
+#define MXNET_OPERATOR_CONTRIB_DYNAMIC_SHAPE_OPS_INL_H_
+
+#include <vector>
+#include "../mxnet_op.h"
+#include "../mshadow_op.h"
+#include "../tensor/matrix_op-inl.h"
+
+namespace mxnet {
+namespace op {
+
+template<typename xpu>
+inline void DynamicReshapeForward(const nnvm::NodeAttrs& attrs,
+ const OpContext &ctx,
+ const std::vector<NDArray> &inputs,
+ const std::vector<OpReqType> &req,
+ const std::vector<NDArray> &outputs) {
+ CHECK_EQ(inputs.size(), 2U);
+ CHECK_EQ(outputs.size(), 1U);
+ const NDArray &out = outputs[0];
+ const NDArray &idx = inputs[1];
+ size_t idx_size = idx.shape()[0];
+ mxnet::TShape shape_value = mxnet::TShape(idx_size, 0);
+ std::vector<index_t> shapev(idx_size, 0);
+
+ // Copy the target shape that is provided in inputs[1]
+ // to the vector shapev
+ MSHADOW_TYPE_SWITCH(idx.dtype(), DType, {
+ DType* idx_dptr = idx.data().dptr<DType>();
+ for (size_t i = 0; i < idx_size; i++) {
+ shapev[i] = static_cast<index_t>(idx_dptr[i]);
+ }
+ });
+ shape_value = InferReshapeShape(mxnet::Tuple<index_t>(shapev),
inputs[0].shape(), false);
+ const_cast<NDArray &>(out).Init(shape_value);
+ mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
+
+ MSHADOW_TYPE_SWITCH(out.dtype(), DType, {
+ MXNET_ASSIGN_REQ_SWITCH(req[0], Req, {
+ mxnet_op::Kernel<mxnet_op::op_with_req<mshadow_op::identity, Req>,
xpu>::Launch(
+ s, inputs[0].data().Size(), out.data().dptr<DType>(),
+ inputs[0].data().dptr<DType>());
+ });
+ });
+}
+
+template<typename xpu>
+inline void DynamicReshapeBackward(const nnvm::NodeAttrs& attrs,
+ const OpContext &ctx,
+ const std::vector<NDArray> &inputs,
+ const std::vector<OpReqType> &req,
+ const std::vector<NDArray> &outputs) {
+ CHECK_EQ(inputs.size(), 1U);
+ CHECK_EQ(outputs.size(), 2U);
+ mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
+
+ MSHADOW_TYPE_SWITCH(outputs[0].dtype(), DType, {
+ MXNET_ASSIGN_REQ_SWITCH(req[0], Req, {
+ mxnet_op::Kernel<mxnet_op::op_with_req<mshadow_op::identity, Req>,
xpu>::Launch(
+ s, inputs[0].data().Size(), outputs[0].data().dptr<DType>(),
+ inputs[0].data().dptr<DType>());
+ });
+ });
+}
+
+} // namespace op
+} // namespace mxnet
+
+#endif // MXNET_OPERATOR_CONTRIB_DYNAMIC_SHAPE_OPS_INL_H_
diff --git a/src/operator/contrib/dynamic_shape_ops.cc
b/src/operator/contrib/dynamic_shape_ops.cc
new file mode 100644
index 0000000..ff71d4b
--- /dev/null
+++ b/src/operator/contrib/dynamic_shape_ops.cc
@@ -0,0 +1,138 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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.
+ */
+/*!
+ * Copyright (c) 2018 by Contributors
+ * \file dynamic_shape_ops.cc
+*/
+
+#include "./dynamic_shape_ops-inl.h"
+#include "../tensor/elemwise_binary_op.h"
+#include "../elemwise_op_common.h"
+
+namespace mxnet {
+namespace op {
+
+inline bool DynamicReshapeType(const nnvm::NodeAttrs& attrs,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ CHECK_EQ(in_attrs->size(), 2U);
+ CHECK_EQ(out_attrs->size(), 1U);
+ TYPE_ASSIGN_CHECK(*out_attrs, 0, (*in_attrs)[0]);
+ TYPE_ASSIGN_CHECK(*in_attrs, 0, (*out_attrs)[0]);
+ return true;
+}
+
+bool DynamicReshapeStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ CHECK_EQ(in_attrs->size(), 2);
+ CHECK_EQ(out_attrs->size(), 1);
+ for (size_t i = 0; i < in_attrs->size(); ++i) {
+ STORAGE_TYPE_ASSIGN_CHECK(*in_attrs, i, kDefaultStorage);
+ }
+ for (size_t i = 0; i < out_attrs->size(); ++i) {
+ STORAGE_TYPE_ASSIGN_CHECK(*out_attrs, i, kDefaultStorage);
+ }
+ DISPATCH_MODE_ASSIGN_CHECK(dispatch_mode, 0, DispatchMode::kFComputeEx);
+ return true;
+}
+
+bool DynamicReshapeBackwardStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ CHECK_EQ(in_attrs->size(), 1);
+ CHECK_EQ(out_attrs->size(), 2);
+ for (size_t i = 0; i < in_attrs->size(); ++i) {
+ STORAGE_TYPE_ASSIGN_CHECK(*in_attrs, i, kDefaultStorage);
+ }
+ for (size_t i = 0; i < out_attrs->size(); ++i) {
+ STORAGE_TYPE_ASSIGN_CHECK(*out_attrs, i, kDefaultStorage);
+ }
+ DISPATCH_MODE_ASSIGN_CHECK(dispatch_mode, 0, DispatchMode::kFComputeEx);
+ return true;
+}
+
+NNVM_REGISTER_OP(_contrib_dynamic_reshape)
+.describe(R"code(
+Experimental support for reshape operator with dynamic shape.
+
+Accepts 2 inputs - data and shape.
+The output returns data in the new shape.
+
+Some dimensions of the shape can take special values from the set {0, -1, -2,
-3, -4}. The significance of each is explained below:
+- ``0`` copy this dimension from the input to the output shape.
+ Example::
+ - input shape = (2,3,4), shape = (4,0,2), output shape = (4,3,2)
+ - input shape = (2,3,4), shape = (2,0,0), output shape = (2,3,4)
+- ``-1`` infers the dimension of the output shape by using the remainder of
the input dimensions
+ keeping the size of the new array same as that of the input array.
+ At most one dimension of shape can be -1.
+ Example::
+ - input shape = (2,3,4), shape = (6,1,-1), output shape = (6,1,4)
+ - input shape = (2,3,4), shape = (3,-1,8), output shape = (3,1,8)
+ - input shape = (2,3,4), shape=(-1,), output shape = (24,)
+- ``-2`` copy all/remainder of the input dimensions to the output shape.
+ Example::
+ - input shape = (2,3,4), shape = (-2,), output shape = (2,3,4)
+ - input shape = (2,3,4), shape = (2,-2), output shape = (2,3,4)
+ - input shape = (2,3,4), shape = (-2,1,1), output shape = (2,3,4,1,1)
+- ``-3`` use the product of two consecutive dimensions of the input shape as
the output dimension.
+ Example::
+ - input shape = (2,3,4), shape = (-3,4), output shape = (6,4)
+ - input shape = (2,3,4,5), shape = (-3,-3), output shape = (6,20)
+ - input shape = (2,3,4), shape = (0,-3), output shape = (2,12)
+ - input shape = (2,3,4), shape = (-3,-2), output shape = (6,4)
+- ``-4`` split one dimension of the input into two dimensions passed
subsequent to -4 in shape (can contain -1).
+ Example::
+ - input shape = (2,3,4), shape = (-4,1,2,-2), output shape =(1,2,3,4)
+ - input shape = (2,3,4), shape = (2,-4,-1,3,-2), output shape = (2,1,3,4)
+
+Example::
+ data = mx.nd.array(np.random.normal(0,1,(2,3,5,5)))
+ shape = mx.nd.array((0,-1))
+ out = mx.sym.contrib.dynamic_reshape(data = data, shape = shape)
+ // out will be of shape (2,75)
+)code" ADD_FILELINE)
+.set_num_inputs(2)
+.set_num_outputs(1)
+.set_attr<nnvm::FListInputNames>("FListInputNames",
+ [](const NodeAttrs& attrs) {
+ return std::vector<std::string>{"data", "shape"};
+ })
+.set_attr<nnvm::FInferType>("FInferType", DynamicReshapeType)
+.set_attr<FInferStorageType>("FInferStorageType", DynamicReshapeStorageType)
+.set_attr<FComputeEx>("FComputeEx<cpu>", DynamicReshapeForward<cpu>)
+.set_attr<nnvm::FGradient>("FGradient",
ElemwiseGradUseNone{"_backward_contrib_dynamic_reshape"})
+.add_argument("data", "NDArray-or-Symbol", "Data")
+.add_argument("shape", "NDArray-or-Symbol", "Shape");
+
+
+NNVM_REGISTER_OP(_backward_contrib_dynamic_reshape)
+.set_num_inputs(1)
+.set_num_outputs(2)
+.set_attr<nnvm::TIsBackward>("TIsBackward", true)
+.set_attr<FInferStorageType>("FInferStorageType",
DynamicReshapeBackwardStorageType)
+.set_attr<FComputeEx>("FComputeEx<cpu>", DynamicReshapeBackward<cpu>);
+
+} // namespace op
+} // namespace mxnet
diff --git a/tests/python/unittest/test_contrib_operator.py
b/tests/python/unittest/test_contrib_operator.py
index 23cfdf8..e3f59de 100644
--- a/tests/python/unittest/test_contrib_operator.py
+++ b/tests/python/unittest/test_contrib_operator.py
@@ -410,3 +410,53 @@ def test_op_mrcnn_mask_target():
assert_almost_equal(mask_targets.asnumpy(), gt_mask_targets.asnumpy())
assert_almost_equal(mask_cls.asnumpy(), gt_mask_cls.asnumpy())
+
+@with_seed()
+def test_dynamic_reshape():
+ def dynamic_reshape_testcases(src_shape, shape_arg, dst_shape):
+ data = mx.sym.Variable('data')
+ shape = mx.sym.Variable('shape')
+ net = mx.sym.contrib.dynamic_reshape(data, shape)
+ js = net.tojson()
+ net = mx.sym.load_json(js)
+ dat_npy = np.random.rand(*src_shape)
+ grad_npy = np.random.rand(*dst_shape)
+ args = {
+ 'data': mx.nd.array(dat_npy),
+ 'shape': mx.nd.array(shape_arg)
+ }
+ args_grad = {
+ 'data': mx.nd.empty(src_shape)
+ }
+ exe = net._bind(default_context(), args, args_grad)
+ exe.forward(is_train=True)
+ assert np.square(exe.outputs[0].asnumpy() -
dat_npy.reshape(dst_shape)).mean() < 1E-7
+ exe.backward(out_grads=mx.nd.array(grad_npy))
+ assert np.square(exe.grad_dict['data'].asnumpy() -
grad_npy.reshape(src_shape)).mean() < 1E-7
+
+ # test ndarray
+ X = mx.nd.random.uniform(shape=src_shape)
+ Y = mx.contrib.nd.dynamic_reshape(X, mx.nd.array(shape_arg))
+ assert_array_equal(Y.shape, dst_shape)
+
+ test_cases = [
+ [(2, 3, 5, 5), (0, -1), (2, 75)],
+ [(2, 3, 5, 5), (0, 0, -1), (2, 3, 25)],
+ [(5, 3, 4, 5), (0, -1, 0), (5, 15, 4)],
+ [(2, 3, 5, 4), (-1, 0, 0), (8, 3, 5)],
+ [(2, 3, 5, 5), (0, 0, 0, 0), (2, 3, 5, 5)],
+ [(2, 4, 5, 3), (-1, 2, 2, 1), (30, 2, 2, 1)],
+ [(2, 3, 5, 6), (-2,), (2, 3, 5, 6)],
+ [(2, 3, 5, 6), (6, 1, -2), (6, 1, 5, 6)],
+ [(2, 3, 5, 6), (-3, -3), (6, 30)],
+ [(2, 3, 5, 6), (-3, -1), (6, 30)],
+ [(64,), (-4, 16, 4), (16, 4)],
+ [(64,), (-4, 16, -1), (16, 4)],
+ [(64, 1, 2, 3), (-4, 16, -1, -2), (16, 4, 1, 2, 3)]]
+
+ for test_case in test_cases:
+ dynamic_reshape_testcases(*test_case)
+
+if __name__ == '__main__':
+ import nose
+ nose.runmodule()
