From 478705fb40463abe20b37af3b9a244f6449b6c95 Mon Sep 17 00:00:00 2001 From: Denisa Roberts Date: Sat, 15 Feb 2020 12:54:30 -0500 Subject: [PATCH] Implement np.random.pareto backward --- python/mxnet/ndarray/numpy/random.py | 8 ++- python/mxnet/numpy/random.py | 4 +- python/mxnet/symbol/numpy/random.py | 8 ++- src/operator/numpy/random/np_pareto_op.cc | 38 +++++++++- src/operator/numpy/random/np_pareto_op.cu | 3 + src/operator/numpy/random/np_pareto_op.h | 88 +++++++++++++++++++---- tests/python/unittest/test_numpy_op.py | 37 +++++++++- 7 files changed, 159 insertions(+), 27 deletions(-) diff --git a/python/mxnet/ndarray/numpy/random.py b/python/mxnet/ndarray/numpy/random.py index 8d99bb1ab1c8..d2667f9bd33a 100644 --- a/python/mxnet/ndarray/numpy/random.py +++ b/python/mxnet/ndarray/numpy/random.py @@ -625,7 +625,7 @@ def weibull(a, size=None, ctx=None, out=None): return _npi.weibull(a=a, size=size, ctx=ctx, out=out) -def pareto(a, size=None): +def pareto(a, size=None, ctx=None, out=None): r"""Draw samples from a Pareto II or Lomax distribution with specified shape a. Parameters @@ -659,13 +659,15 @@ def pareto(a, size=None): """ from ...numpy import ndarray as np_ndarray tensor_type_name = np_ndarray + if ctx is None: + ctx = current_context() if size == (): size = None is_tensor = isinstance(a, tensor_type_name) if is_tensor: - return _npi.pareto(a, a=None, size=size) + return _npi.pareto(a, a=None, size=size, ctx=ctx, out=out) else: - return _npi.pareto(a=a, size=size) + return _npi.pareto(a=a, size=size, ctx=ctx, out=out) def power(a, size=None): diff --git a/python/mxnet/numpy/random.py b/python/mxnet/numpy/random.py index 3c05ff17de0f..21bafe0deef8 100644 --- a/python/mxnet/numpy/random.py +++ b/python/mxnet/numpy/random.py @@ -656,7 +656,7 @@ def weibull(a, size=None, ctx=None, out=None): return _mx_nd_np.random.weibull(a, size=size, ctx=ctx, out=out) -def pareto(a, size=None): +def pareto(a, size=None, ctx=None, out=None): r"""Draw samples from a Pareto II or Lomax distribution with specified shape a. Parameters @@ -688,7 +688,7 @@ def pareto(a, size=None): where a is the shape and m the scale. Here m is assumed 1. The Pareto distribution is a power law distribution. Pareto created it to describe the wealth in the economy. """ - return _mx_nd_np.random.pareto(a, size) + return _mx_nd_np.random.pareto(a, size=size, ctx=ctx, out=out) def power(a, size=None): diff --git a/python/mxnet/symbol/numpy/random.py b/python/mxnet/symbol/numpy/random.py index 5885488550f5..697025a5cd62 100644 --- a/python/mxnet/symbol/numpy/random.py +++ b/python/mxnet/symbol/numpy/random.py @@ -695,7 +695,7 @@ def weibull(a, size=None, ctx=None, out=None): return _npi.weibull(a=a, size=size, ctx=ctx, out=out) -def pareto(a, size=None): +def pareto(a, size=None, ctx=None, out=None): r"""Draw samples from a Pareto II or Lomax distribution with specified shape a. Parameters @@ -729,13 +729,15 @@ def pareto(a, size=None): """ from ..numpy import _Symbol as np_symbol tensor_type_name = np_symbol + if ctx is None: + ctx = current_context() if size == (): size = None is_tensor = isinstance(a, tensor_type_name) if is_tensor: - return _npi.pareto(a, a=None, size=size) + return _npi.pareto(a, a=None, size=size, ctx=ctx, out=out) else: - return _npi.pareto(a=a, size=size) + return _npi.pareto(a=a, size=size, ctx=ctx, out=out) def power(a, size=None): diff --git a/src/operator/numpy/random/np_pareto_op.cc b/src/operator/numpy/random/np_pareto_op.cc index df77448907fe..bdfb8ca8a87b 100644 --- a/src/operator/numpy/random/np_pareto_op.cc +++ b/src/operator/numpy/random/np_pareto_op.cc @@ -32,6 +32,7 @@ namespace op { DMLC_REGISTER_PARAMETER(NumpyParetoParam); NNVM_REGISTER_OP(_npi_pareto) +.describe("Numpy behavior Pareto") .set_num_inputs( [](const nnvm::NodeAttrs& attrs) { const NumpyParetoParam& param = nnvm::get(attrs.parsed); @@ -41,7 +42,11 @@ NNVM_REGISTER_OP(_npi_pareto) } return num_inputs; }) -.set_num_outputs(1) +.set_num_outputs(2) +.set_attr("FNumVisibleOutputs", + [](const NodeAttrs& attrs){ + return 1; + }) .set_attr("FListInputNames", [](const NodeAttrs& attrs) { const NumpyParetoParam& param = nnvm::get(attrs.parsed); @@ -52,10 +57,11 @@ NNVM_REGISTER_OP(_npi_pareto) return (num_inputs == 0) ? std::vector() : std::vector{"input1"}; }) .set_attr_parser(ParamParser) -.set_attr("FInferShape", UnaryDistOpShape) +.set_attr("FInferShape", TwoparamsDistOpShape) .set_attr("FInferType", [](const nnvm::NodeAttrs &attrs, std::vector *in_attrs, std::vector *out_attrs) { (*out_attrs)[0] = mshadow::kFloat32; + (*out_attrs)[1] = mshadow::kFloat32; return true; }) .set_attr("FResourceRequest", @@ -64,9 +70,35 @@ NNVM_REGISTER_OP(_npi_pareto) ResourceRequest::kRandom, ResourceRequest::kTempSpace}; }) .set_attr("FCompute", NumpyParetoForward) -.set_attr("FGradient", MakeZeroGradNodes) +.set_attr("FGradient", ElemwiseGradUseInOut{"_backward_broadcast_pareto"}) .add_argument("input1", "NDArray-or-Symbol", "Source input") .add_arguments(NumpyParetoParam::__FIELDS__()); +NNVM_REGISTER_OP(_backward_broadcast_pareto) +.set_attr("TIsBackward", true) +.set_attr_parser(ParamParser) +.set_num_inputs( + [](const nnvm::NodeAttrs& attrs){ + const NumpyParetoParam& param = nnvm::get(attrs.parsed); + int num_inputs = 5; + if (param.a.has_value()) num_inputs -= 1; + return num_inputs; + } +) + .set_num_outputs( + [](const nnvm::NodeAttrs& attrs){ + const NumpyParetoParam& param = nnvm::get(attrs.parsed); + int num_outputs = 1; + if (param.a.has_value()) num_outputs -= 1; + return num_outputs; + } + ) + .set_attr("FResourceRequest", + [](const NodeAttrs& attrs){ + return std::vector{ResourceRequest::kTempSpace}; + }) + .set_attr("FCompute", ParetoReparamBackward) + .add_arguments(NumpyParetoParam::__FIELDS__()); + } // namespace op } // namespace mxnet diff --git a/src/operator/numpy/random/np_pareto_op.cu b/src/operator/numpy/random/np_pareto_op.cu index 9af362cc69ce..d8a8a896e653 100644 --- a/src/operator/numpy/random/np_pareto_op.cu +++ b/src/operator/numpy/random/np_pareto_op.cu @@ -31,5 +31,8 @@ namespace op { NNVM_REGISTER_OP(_npi_pareto) .set_attr("FCompute", NumpyParetoForward); +NNVM_REGISTER_OP(_backward_broadcast_pareto) +.set_attr("FCompute", ParetoReparamBackward); + } // namespace op } // namespace mxnet diff --git a/src/operator/numpy/random/np_pareto_op.h b/src/operator/numpy/random/np_pareto_op.h index 01bf29a6a8b0..85eab97aef8c 100644 --- a/src/operator/numpy/random/np_pareto_op.h +++ b/src/operator/numpy/random/np_pareto_op.h @@ -44,6 +44,7 @@ namespace op { struct NumpyParetoParam : public dmlc::Parameter { dmlc::optional a; dmlc::optional> size; + std::string ctx; DMLC_DECLARE_PARAMETER(NumpyParetoParam) { DMLC_DECLARE_FIELD(a) .set_default(dmlc::optional()); @@ -52,14 +53,17 @@ struct NumpyParetoParam : public dmlc::Parameter { .describe("Output shape. If the given shape is, " "e.g., (m, n, k), then m * n * k samples are drawn. " "Default is None, in which case a single value is returned."); + DMLC_DECLARE_FIELD(ctx).set_default("cpu").describe( + "Context of output, in format [cpu|gpu|cpu_pinned](n)." + " Only used for imperative calls."); } }; template struct scalar_pareto_kernel { - MSHADOW_XINLINE static void Map(index_t i, float a, float *threshold, + MSHADOW_XINLINE static void Map(index_t i, float a, float *noise, DType *out) { - out[i] = exp(-log(threshold[i])/a) - DType(1); + out[i] = exp(-log(noise[i])/a) - DType(1); } }; @@ -67,7 +71,7 @@ namespace mxnet_op { template struct check_legal_a_kernel { - MSHADOW_XINLINE static void Map(index_t i, IType *a, float* flag) { + MSHADOW_XINLINE static void Map(index_t i, IType *a, float *flag) { if (a[i] <= 0.0) { flag[0] = -1.0; } @@ -80,10 +84,13 @@ struct pareto_kernel { MSHADOW_XINLINE static void Map(index_t i, const Shape &stride, const Shape &oshape, - IType *aparams, float* threshold, OType *out) { + IType *aparams, float *noise, OType *out) { Shape coord = unravel(i, oshape); auto idx = static_cast(dot(coord, stride)); - out[i] = exp(-log(threshold[i])/aparams[idx]) - IType(1); + noise[i] = -log(noise[i]); + out[i] = exp(noise[i]/aparams[idx]) - IType(1); + // get grad + noise[i] = -noise[i] * (out[i] + 1.0) * (1.0/(aparams[idx] * aparams[idx])); } }; @@ -91,24 +98,23 @@ struct pareto_kernel { template void NumpyParetoForward(const nnvm::NodeAttrs &attrs, - const OpContext &ctx, - const std::vector &inputs, - const std::vector &req, - const std::vector &outputs) { + const OpContext &ctx, + const std::vector &inputs, + const std::vector &req, + const std::vector &outputs) { using namespace mshadow; using namespace mxnet_op; const NumpyParetoParam ¶m = nnvm::get(attrs.parsed); Stream *s = ctx.get_stream(); - index_t output_len = outputs[0].Size(); Random *prnd = ctx.requested[0].get_random(s); Tensor workspace = - ctx.requested[1].get_space_typed(Shape1(output_len + 1), s); - Tensor uniform_tensor = workspace.Slice(0, output_len); - Tensor indicator_device = workspace.Slice(output_len, output_len + 1); + ctx.requested[1].get_space_typed(Shape1(1), s); + Tensor uniform_tensor = outputs[1].FlatTo1D(s); + Tensor indicator_device = workspace; float indicator_host = 1.0; float *indicator_device_ptr = indicator_device.dptr_; Kernel::Launch(s, 1, indicator_device_ptr); - prnd->SampleUniform(&workspace, 0.0, 1.0); + prnd->SampleUniform(&uniform_tensor, 0.0, 1.0); if (param.a.has_value()) { CHECK_GT(param.a.value(), 0.0) << "ValueError: expect a > 0"; MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, DType, { @@ -140,6 +146,60 @@ void NumpyParetoForward(const nnvm::NodeAttrs &attrs, } } +template +inline void ScalarParetoReparamBackwardImpl(const OpContext& ctx, + const std::vector& inputs, + const std::vector& req, + const std::vector& outputs, + const mxnet::TShape& new_ishape, + const mxnet::TShape& new_oshape) { + using namespace mshadow; + using namespace mshadow::expr; + using namespace broadcast; + Stream *s = ctx.get_stream(); + const TBlob igrad = outputs[0].reshape(new_ishape); + // inputs: [grad_from_samples, grad_from_noise(invisible), input_tensor, + // samples, noise] + const TBlob ograd = inputs[0].reshape(new_oshape); + const TBlob itensor = inputs[2].reshape(new_ishape); + const TBlob samples = inputs[3].reshape(new_oshape); + const TBlob noise = inputs[4].reshape(new_oshape); + size_t workspace_size = + ReduceWorkspaceSize(s, igrad.shape_, req[0], ograd.shape_); + Tensor workspace = + ctx.requested[0].get_space_typed(Shape1(workspace_size), s); + Reduce( + s, igrad, req[0], workspace, ograd, noise, noise); + } + +template +void ParetoReparamBackward(const nnvm::NodeAttrs& attrs, + const OpContext& ctx, + const std::vector& inputs, + const std::vector& reqs, + const std::vector& outputs) { +// skip kernel launch for zero-size tensors +if (inputs[0].shape_.Size() == 0U) { + return; +} +// [scalar] case +if (outputs.size() == 0U) { + return; +} +// [tensor] case +if (inputs.size() == 5U) { + mxnet::TShape new_ishape, new_oshape; + int ndim = FillShape(outputs[0].shape_, outputs[0].shape_, inputs[0].shape_, + &new_ishape, &new_ishape, &new_oshape); + MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, DType, { + BROADCAST_NDIM_SWITCH(ndim, NDim, { + ScalarParetoReparamBackwardImpl( + ctx, inputs, reqs, outputs, new_ishape, new_oshape); + }); + }); + } +} + } // namespace op } // namespace mxnet diff --git a/tests/python/unittest/test_numpy_op.py b/tests/python/unittest/test_numpy_op.py index 759732e27aea..d67f4f2bc769 100644 --- a/tests/python/unittest/test_numpy_op.py +++ b/tests/python/unittest/test_numpy_op.py @@ -4035,14 +4035,13 @@ def hybrid_forward(self, F, a): expected_shape = a.shape assert mx_out.shape == expected_shape - # test illegal parameter values (as numpy produces) + # test illegal parameter values def _test_exception(a): output = op(a=a).asnumpy() for op in op_names: op = getattr(np.random, op_name, None) if op is not None: assertRaises(ValueError, _test_exception, -1) - if op in ['pareto', 'power']: assertRaises(ValueError, _test_exception, 0) @@ -4079,6 +4078,40 @@ def hybrid_forward(self, F, a): assert_almost_equal(a.grad.asnumpy().sum(), formula_grad.asnumpy().sum(), rtol=1e-3, atol=1e-5) +@with_seed() +@use_np +def test_np_pareto_grad(): + class TestRandomP(HybridBlock): + def __init__(self, shape): + super(TestRandomP, self).__init__() + self._shape = shape + + def hybrid_forward(self, F, a): + return F.np.random.pareto(a, self._shape) + + output_shapes = [ + (3, 2), + (4, 3, 2, 2), + (3, 4, 5) + ] + for hybridize in [False, True]: + for out_shape in output_shapes: + test_w_grad = TestRandomP(out_shape) + if hybridize: + test_w_grad.hybridize() + a = np.ones(out_shape) + a.attach_grad() + with mx.autograd.record(): + mx_out = test_w_grad(a) + mx_out.backward() + + # gradient formula from calculus (a=1) + noise = np.log(mx_out + np.ones(mx_out.shape)) + formula_grad = - (mx_out + np.ones(mx_out.shape)) * noise + assert a.grad.shape == out_shape + assert_almost_equal(a.grad.asnumpy().sum(), formula_grad.asnumpy().sum(), rtol=1e-3, atol=1e-5) + + @with_seed() @use_np def test_np_randn():