diff --git a/python/mxnet/ndarray/numpy/_op.py b/python/mxnet/ndarray/numpy/_op.py index c5d9d7fc206a..e6557e9ce145 100644 --- a/python/mxnet/ndarray/numpy/_op.py +++ b/python/mxnet/ndarray/numpy/_op.py @@ -33,7 +33,7 @@ 'rint', 'radians', 'reciprocal', 'square', 'negative', 'fix', 'ceil', 'floor', 'trunc', 'logical_not', 'arcsinh', 'arccosh', 'arctanh', 'tensordot', 'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', 'stack', 'mean', - 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax'] + 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var'] @set_module('mxnet.ndarray.numpy') @@ -2201,3 +2201,17 @@ def mean(a, axis=None, dtype=None, out=None, keepdims=False): # pylint: disable array(0.55) """ return _npi.mean(a, axis=axis, dtype=dtype, keepdims=keepdims, out=out) + + +@set_module('mxnet.ndarray.numpy') +def std(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False): + """ + """ + return _npi.std(a, axis=axis, dtype=dtype, ddof=ddof, keepdims=keepdims, out=out) + + +@set_module('mxnet.ndarray.numpy') +def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False): + """ + """ + return _npi.var(a, axis=axis, dtype=dtype, ddof=ddof, keepdims=keepdims, out=out) diff --git a/python/mxnet/numpy/multiarray.py b/python/mxnet/numpy/multiarray.py index 5c0b2082c7de..ca84bcc33ed6 100644 --- a/python/mxnet/numpy/multiarray.py +++ b/python/mxnet/numpy/multiarray.py @@ -52,7 +52,7 @@ 'degrees', 'log2', 'log1p', 'rint', 'radians', 'reciprocal', 'square', 'negative', 'fix', 'ceil', 'floor', 'trunc', 'logical_not', 'arcsinh', 'arccosh', 'arctanh', 'tensordot', 'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', - 'stack', 'mean', 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax'] + 'stack', 'mean', 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var'] # Return code for dispatching indexing function call _NDARRAY_UNSUPPORTED_INDEXING = -1 @@ -1172,11 +1172,9 @@ def mean(self, axis=None, dtype=None, out=None, keepdims=False): # pylint: disa """Returns the average of the array elements along given axis.""" raise NotImplementedError - # TODO(junwu): Use mxnet std op instead of onp.std def std(self, axis=None, dtype=None, out=None, ddof=0, keepdims=False): # pylint: disable=arguments-differ """Returns the standard deviation of the array elements along given axis.""" - ret_np = self.asnumpy().std(axis=axis, dtype=dtype, out=out, ddof=ddof, keepdims=keepdims) - return array(ret_np, dtype=ret_np.dtype, ctx=self.context) + return _mx_np_op.std(self, axis=axis, dtype=dtype, ddof=ddof, keepdims=keepdims, out=out) def cumsum(self, axis=None, dtype=None, out=None): """Return the cumulative sum of the elements along the given axis.""" @@ -3644,3 +3642,17 @@ def mean(a, axis=None, dtype=None, out=None, keepdims=False): # pylint: disable array(0.55) """ return _npi.mean(a, axis=axis, dtype=dtype, keepdims=keepdims, out=out) + + +@set_module('mxnet.numpy') +def std(a, axis=None, dtype=None, out=None, ddof=0, keepdims=None): + """ + """ + return _npi.std(a, axis=axis, dtype=dtype, ddof=ddof, keepdims=keepdims, out=out) + + +@set_module('mxnet.numpy') +def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=None): + """ + """ + return _npi.var(a, axis=axis, dtype=dtype, ddof=ddof, keepdims=keepdims, out=out) diff --git a/python/mxnet/symbol/numpy/_symbol.py b/python/mxnet/symbol/numpy/_symbol.py index e79e3fa5ced0..999b32f8516a 100644 --- a/python/mxnet/symbol/numpy/_symbol.py +++ b/python/mxnet/symbol/numpy/_symbol.py @@ -35,7 +35,7 @@ 'rint', 'radians', 'reciprocal', 'square', 'negative', 'fix', 'ceil', 'floor', 'trunc', 'logical_not', 'arcsinh', 'arccosh', 'arctanh', 'tensordot', 'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', 'stack', 'mean', - 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax'] + 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var'] def _num_outputs(sym): @@ -2569,4 +2569,18 @@ def mean(a, axis=None, dtype=None, out=None, keepdims=False): # pylint: disable return _npi.mean(a, axis=axis, dtype=dtype, keepdims=keepdims, out=out) +@set_module('mxnet.symbol.numpy') +def std(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False): + """ + """ + return _npi.std(a, axis=axis, dtype=dtype, ddof=ddof, keepdims=keepdims, out=out) + + +@set_module('mxnet.symbol.numpy') +def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False): + """ + """ + return _npi.var(a, axis=axis, dtype=dtype, ddof=ddof, keepdims=keepdims, out=out) + + _set_np_symbol_class(_Symbol) diff --git a/src/operator/numpy/np_broadcast_reduce_op.h b/src/operator/numpy/np_broadcast_reduce_op.h index e9d815dec5f1..1388f37b9120 100644 --- a/src/operator/numpy/np_broadcast_reduce_op.h +++ b/src/operator/numpy/np_broadcast_reduce_op.h @@ -27,6 +27,7 @@ #include #include +#include "../nn/moments-inl.h" #include "../tensor/broadcast_reduce_op.h" namespace mxnet { @@ -230,6 +231,139 @@ void NumpyReduceAxesBackwardUseInOut(const nnvm::NodeAttrs& attrs, ReduceAxesBackwardUseInOutImpl(ctx, small, inputs, req, outputs); } +struct NumpyMomentsParam : public dmlc::Parameter { + dmlc::optional> axis; + dmlc::optional dtype; + bool keepdims; + int ddof; + DMLC_DECLARE_PARAMETER(NumpyMomentsParam) { + DMLC_DECLARE_FIELD(axis) + .set_default(dmlc::optional>()) + .describe("Axis or axes along which a sum is performed. The default, axis=None, will sum " + "all of the elements of the input array. If axis is negative it counts from the " + "last to the first axis."); + DMLC_DECLARE_FIELD(dtype) + .add_enum("float16", mshadow::kFloat16) + .add_enum("float32", mshadow::kFloat32) + .add_enum("float64", mshadow::kFloat64) + .add_enum("int8", mshadow::kInt8) + .add_enum("int32", mshadow::kInt32) + .add_enum("int64", mshadow::kInt64) + .set_default(dmlc::optional()) + .describe("The type of the returned array and of the accumulator in which the elements are " + "summed. The dtype of a is used by default unless a has an integer dtype of less " + "precision than the default platform integer. In that case, if a is signed then " + "the platform integer is used while if a is unsigned then an unsigned integer of " + "the same precision as the platform integer is used."); + DMLC_DECLARE_FIELD(ddof).set_default(0) + .describe("Starting value for the sum."); + DMLC_DECLARE_FIELD(keepdims).set_default(false) + .describe("If this is set to `True`, the reduced axes are left " + "in the result as dimension with size one."); + } +}; + +template +void ReduceAxesComputeWithWorkspaceImpl(const OpContext& ctx, + const std::vector& inputs, + const std::vector& req, + const std::vector& outputs, + mshadow::Tensor& workspace, + const mxnet::TShape& src_shape, + const mxnet::TShape& dst_shape, + const int ddof = 0) { + using namespace mshadow; + using namespace mshadow::expr; + + Stream *s = ctx.get_stream(); + MSHADOW_TYPE_SWITCH(inputs[0].type_flag_, DType, { + MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, OType, { + const TBlob in_data = inputs[0].reshape(src_shape); + const TBlob out_data = outputs[0].reshape(dst_shape); + BROADCAST_NDIM_SWITCH(dst_shape.ndim(), NDim, { + broadcast::Reduce( + s, out_data, req[0], workspace, in_data); + if (normalize) { + auto out = out_data.FlatTo2D(s); + out /= scalar(src_shape.Size()/dst_shape.Size() - ddof); + } + }); + }); + }); +} + +template +void NumpyMomentsForward(const nnvm::NodeAttrs& attrs, + const OpContext& ctx, + const std::vector& inputs, + const std::vector& req, + const std::vector& outputs) { + using namespace mshadow; + using namespace mshadow::expr; + using namespace mshadow_op; + using namespace mxnet_op; + + CHECK_EQ(inputs.size(), 1U); + CHECK_EQ(req.size(), 2U); + CHECK_EQ(outputs.size(), 2U); + + const NumpyMomentsParam& param = nnvm::get(attrs.parsed); + + Stream *s = ctx.get_stream(); + + const TBlob& data = inputs[0]; + const TBlob& moment = outputs[0]; + const TBlob& mean = outputs[1]; + + mxnet::TShape small; + if (param.keepdims) { + small = moment.shape_; + } else { + small = NumpyReduceAxesShapeImpl(data.shape_, param.axis, true); + } + + mxnet::TShape src_shape, dst_shape; + BroadcastReduceShapeCompact(data.shape_, small, &src_shape, &dst_shape); + + MSHADOW_TYPE_SWITCH(inputs[0].type_flag_, DType, { + MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, OType, { + // Get workspace and temp space for data - mean + size_t workspace_size = 0; + BROADCAST_NDIM_SWITCH(dst_shape.ndim(), NDim, { + workspace_size = broadcast::ReduceWorkspaceSize( + s, dst_shape, req[0], src_shape);; + }); + size_t temp_data_size = data.shape_.Size() * sizeof(DType); + size_t temp_mem_size = temp_data_size + workspace_size; + Tensor temp_mem = + ctx.requested[0].get_space_typed(Shape1(temp_mem_size), s); + DType *temp_data_ptr = reinterpret_cast(temp_mem.dptr_); + char *workspace_ptr = temp_mem.dptr_ + temp_data_size; + Tensor workspace(workspace_ptr, Shape1(workspace_size), s); + // Compute mean + ReduceAxesComputeWithWorkspaceImpl( + ctx, inputs, {kWriteTo}, {mean}, workspace, src_shape, dst_shape); + // Compute data - mean + Shape<6> data_shape, mean_shape; + for (int i = 0; i < 6; ++i) { + data_shape[i] = (i < data.shape_.ndim()) ? data.shape_[i] : 1; + mean_shape[i] = (i < small.ndim()) ? small[i] : 1; + } + Kernel::Launch(s, data_shape.Size(), temp_data_ptr, + data.dptr(), mean.dptr(), data_shape, mean_shape); + Tensor temp_data_tensor(temp_data_ptr, Shape1(data.shape_.Size()), s); + TBlob temp_data_blob = TBlob(temp_data_tensor).reshape(data.shape_); + ReduceAxesComputeWithWorkspaceImpl( + ctx, {temp_data_blob}, {req[0]}, {moment}, workspace, src_shape, dst_shape, param.ddof); + if (sqrt) { + Tensor moment_tensor = moment.FlatTo1D(s); + moment_tensor = F(moment_tensor); + } + }); + }); +} + template void NumpyBroadcastToForward(const nnvm::NodeAttrs& attrs, const OpContext& ctx, diff --git a/src/operator/numpy/np_broadcast_reduce_op_value.cc b/src/operator/numpy/np_broadcast_reduce_op_value.cc index 1bd75c82ba79..9b79b97bcfd6 100644 --- a/src/operator/numpy/np_broadcast_reduce_op_value.cc +++ b/src/operator/numpy/np_broadcast_reduce_op_value.cc @@ -29,6 +29,7 @@ namespace mxnet { namespace op { DMLC_REGISTER_PARAMETER(NumpyReduceAxesParam); +DMLC_REGISTER_PARAMETER(NumpyMomentsParam); inline bool NumpySumType(const nnvm::NodeAttrs& attrs, std::vector *in_attrs, @@ -153,6 +154,94 @@ NNVM_REGISTER_OP(_backward_np_mean) .set_num_inputs(1) .set_attr("FCompute", NumpyReduceAxesBackwardUseNone); +inline bool NumpyMomentsShape(const nnvm::NodeAttrs& attrs, + std::vector *in_attrs, + std::vector *out_attrs) { + CHECK_EQ(in_attrs->size(), 1U); + CHECK_EQ(out_attrs->size(), 2U); + if (!shape_is_known(in_attrs->at(0))) { + return false; + } + const NumpyMomentsParam& param = nnvm::get(attrs.parsed); + mxnet::TShape out_shape = NumpyReduceAxesShapeImpl((*in_attrs)[0], param.axis, param.keepdims); + SHAPE_ASSIGN_CHECK(*out_attrs, 0, out_shape); + SHAPE_ASSIGN_CHECK(*out_attrs, 1, out_shape); + + return shape_is_known(out_attrs->at(0)) && shape_is_known(out_attrs->at(1)); +} + +inline bool NumpyMomentsType(const nnvm::NodeAttrs& attrs, + std::vector *in_attrs, + std::vector *out_attrs) { + CHECK_EQ(in_attrs->size(), 1U); + CHECK_EQ(out_attrs->size(), 2U); + const NumpyMomentsParam ¶m = nnvm::get(attrs.parsed); + + if (param.dtype.has_value()) { + TYPE_ASSIGN_CHECK(*out_attrs, 0, param.dtype.value()); + } else { + TYPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(0)); + TYPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(0)); + } + TYPE_ASSIGN_CHECK(*out_attrs, 1, in_attrs->at(0)); + + return out_attrs->at(0) != -1 && in_attrs->at(0) != -1; +} + +NNVM_REGISTER_OP(_npi_std) +.set_num_inputs(1) +.set_num_outputs(2) +.set_attr_parser(ParamParser) +.set_attr("FInferShape", NumpyMomentsShape) +.set_attr("FInferType", NumpyMomentsType) +.set_attr("FListInputNames", + [](const NodeAttrs& attrs) { + return std::vector{"a"}; + }) +.set_attr("FListOutputNames", + [](const NodeAttrs& attrs) { + return std::vector{"std", "mean"}; + }) +.set_attr("FNumVisibleOutputs", + [](const NodeAttrs& attrs) { + return 1; + }) +.add_argument("a", "NDArray-or-Symbol", "The input") +.add_arguments(NumpyMomentsParam::__FIELDS__()) +.set_attr("FCompute", NumpyMomentsForward) +.set_attr("FResourceRequest", + [](const NodeAttrs& attrs) { + return std::vector{ResourceRequest::kTempSpace}; + }) +.set_attr("FGradient", MakeZeroGradNodes); + +NNVM_REGISTER_OP(_npi_var) +.set_num_inputs(1) +.set_num_outputs(2) +.set_attr_parser(ParamParser) +.set_attr("FInferShape", NumpyMomentsShape) +.set_attr("FInferType", NumpyMomentsType) +.set_attr("FListInputNames", + [](const NodeAttrs& attrs) { + return std::vector{"a"}; + }) +.set_attr("FListOutputNames", + [](const NodeAttrs& attrs) { + return std::vector{"var", "mean"}; + }) +.set_attr("FNumVisibleOutputs", + [](const NodeAttrs& attrs) { + return 1; + }) +.add_argument("a", "NDArray-or-Symbol", "The input") +.add_arguments(NumpyMomentsParam::__FIELDS__()) +.set_attr("FCompute", NumpyMomentsForward) +.set_attr("FResourceRequest", + [](const NodeAttrs& attrs) { + return std::vector{ResourceRequest::kTempSpace}; + }) +.set_attr("FGradient", MakeZeroGradNodes); + bool NumpyBroadcastToShape(const nnvm::NodeAttrs& attrs, mxnet::ShapeVector *in_attrs, mxnet::ShapeVector *out_attrs) { diff --git a/src/operator/numpy/np_broadcast_reduce_op_value.cu b/src/operator/numpy/np_broadcast_reduce_op_value.cu index 71b267c102d7..b165a198c823 100644 --- a/src/operator/numpy/np_broadcast_reduce_op_value.cu +++ b/src/operator/numpy/np_broadcast_reduce_op_value.cu @@ -44,6 +44,12 @@ NNVM_REGISTER_OP(_npi_mean) NNVM_REGISTER_OP(_backward_np_mean) .set_attr("FCompute", NumpyReduceAxesBackwardUseNone); +NNVM_REGISTER_OP(_npi_std) +.set_attr("FCompute", NumpyMomentsForward); + +NNVM_REGISTER_OP(_npi_var) +.set_attr("FCompute", NumpyMomentsForward); + NNVM_REGISTER_OP(_np_broadcast_to) .set_attr("FCompute", NumpyBroadcastToForward); diff --git a/tests/python/unittest/test_numpy_op.py b/tests/python/unittest/test_numpy_op.py index a2eb12ca93d6..171edbf79b80 100644 --- a/tests/python/unittest/test_numpy_op.py +++ b/tests/python/unittest/test_numpy_op.py @@ -341,6 +341,70 @@ def is_int(dtype): assert_almost_equal(mx_out.asnumpy(), np_out, rtol=1e-3, atol=1e-5) +@with_seed() +@use_np +def test_np_moment(): + class TestMoment(HybridBlock): + def __init__(self, name, axis=None, dtype=None, keepdims=False, ddof=0): + super(TestMoment, self).__init__() + self._name = name + self._axis = axis + self._dtype = dtype + self._keepdims = keepdims + self._ddof = ddof + + def hybrid_forward(self, F, a, *args, **kwargs): + return getattr(F.np, self._name)(a, axis=self._axis, dtype=self._dtype, keepdims=self._keepdims, ddof=self._ddof) + + def is_int(dtype): + return 'int' in dtype + + def legalize_shape(shape): + shape_ = list(shape) + for i in range(len(shape_)): + shape_[i] += 1 + return tuple(shape_) + + in_data_dim = random.choice([2, 3, 4]) + shape = rand_shape_nd(in_data_dim, dim=3) + shape = legalize_shape(shape) + acc_type = {'float16': 'float32', 'float32': 'float64', 'float64': 'float64', + 'int8': 'float64', 'int32': 'float64', 'int64': 'float64'} + + for name in ['var', 'std']: + for hybridize in [False, True]: + for ddof in [0, 1]: + for keepdims in [True, False]: + for axis in ([i for i in range(in_data_dim)] + [(), None]): + for itype in ['float16', 'float32', 'float64', 'int8', 'int32', 'int64']: + for dtype in ['float16', 'float32', 'float64']: + if is_int(dtype) and not is_int(itype) or is_int(itype) and is_int(dtype): + continue + atol = 1e-4 if itype == 'float16' or dtype == 'float16' else 1e-5 + rtol = 1e-2 if itype == 'float16' or dtype == 'float16' else 1e-3 + # test gluon + test_moment = TestMoment(name, axis=axis, dtype=dtype, keepdims=keepdims, ddof=ddof) + if hybridize: + test_moment.hybridize() + if is_int(itype): + x = _np.random.randint(-16, 16, shape, dtype=itype) + x = mx.nd.array(x) + else: + x = mx.nd.random.uniform(-1.0, 1.0, shape=shape, dtype=itype) + x = x.as_np_ndarray() + x.attach_grad() + expected_ret = getattr(_np, name)(x.asnumpy(), axis=axis, dtype=acc_type[itype], keepdims=keepdims, ddof=ddof) + expected_ret = expected_ret.astype(dtype) + y = test_moment(x) + assert y.shape == expected_ret.shape + assert_almost_equal(y.asnumpy(), expected_ret, rtol=rtol, atol=atol, use_broadcast=False, equal_nan=True) + + # test imperative + mx_out = getattr(np, name)(x, axis=axis, dtype=dtype, keepdims=keepdims, ddof=ddof) + np_out = getattr(_np, name)(x.asnumpy(), axis=axis, dtype=acc_type[itype], keepdims=keepdims, ddof=ddof).astype(dtype) + assert_almost_equal(mx_out.asnumpy(), np_out, rtol=rtol, atol=atol, use_broadcast=False, equal_nan=True) + + @with_seed() @use_np def test_np_linspace():