From 0a08a3bf77d9583e82867169c0611f09beab597b Mon Sep 17 00:00:00 2001 From: kshitij12345 Date: Sat, 29 Jun 2019 12:53:22 +0530 Subject: [PATCH 1/3] add higher order support for reciprocal and abs --- .../tensor/elemwise_unary_op_basic.cc | 54 ++++++++++++++++++- 1 file changed, 52 insertions(+), 2 deletions(-) diff --git a/src/operator/tensor/elemwise_unary_op_basic.cc b/src/operator/tensor/elemwise_unary_op_basic.cc index 98dc8dad825f..29e45ce396ad 100644 --- a/src/operator/tensor/elemwise_unary_op_basic.cc +++ b/src/operator/tensor/elemwise_unary_op_basic.cc @@ -689,7 +689,38 @@ Example:: MXNET_OPERATOR_REGISTER_BINARY(_backward_reciprocal) .set_attr("FCompute", - ElemwiseBinaryOp::Compute >); + ElemwiseBinaryOp::Compute >) +.set_attr("FGradient", + [](const nnvm::NodePtr& n, const std::vector& ograds) { + // ograds[0]: dL/dxgrad + // inputs[0]: dL/dy + // inputs[1]: x + // f(x) = y = 1/x + // f'(x) = -1/x^2 + // f''(x) = 2/x^3 = -2 * (f'(x) * f(x)) + + const std::unordered_map args = {{"scalar", "-2.0"}}; + + auto dydx_mul_dldy = nnvm::NodeEntry{n}; // f'(x) * head_grads + auto dydx = MakeNode("elemwise_div", n->attrs.name + "_dydx", + {dydx_mul_dldy, n->inputs[0]}, nullptr, &n); + auto fx = MakeNode("reciprocal", n->attrs.name + "_fx", + {n->inputs[1]}, nullptr, &n); + + auto d2ydx2_mid = MakeNode("elemwise_mul", n->attrs.name + "_d2ydx2_mid", + {dydx_mul_dldy, nnvm::NodeEntry{fx}}, nullptr, &n); + + auto d2ydx2 = MakeNode("_mul_scalar", n->attrs.name + "_d2ydx2", + {nnvm::NodeEntry{d2ydx2_mid}}, &args, &n); + + std::vector ret; + + ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad", + {ograds[0], nnvm::NodeEntry{dydx}}, nullptr, &n)); + ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad_inp", + {ograds[0], nnvm::NodeEntry{d2ydx2}}, nullptr, &n)); + return ret; +}); // abs MXNET_OPERATOR_REGISTER_UNARY_WITH_RSP_CSR(abs, cpu, mshadow_op::abs) @@ -708,7 +739,26 @@ The storage type of ``abs`` output depends upon the input storage type: )code" ADD_FILELINE) .set_attr("FGradient", ElemwiseGradUseIn{"_backward_abs"}); -MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU(_backward_abs, unary_bwd); +MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU(_backward_abs, unary_bwd) +.set_attr("FGradient", + [](const nnvm::NodePtr& n, const std::vector& ograds) { + // ograds[0]: dL/dxgrad + // inputs[0]: dL/dy + // inputs[1]: y + // f(x) -> abs(x) + // f'(x) = 1 if x > 0 else -1 + // f''(x) = 0 + auto dydx = MakeNode("elemwise_div", n->attrs.name + "_dydx", + {nnvm::NodeEntry{n}, n->inputs[0]}, nullptr, &n); + + std::vector ret; + ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad", + {ograds[0], nnvm::NodeEntry(dydx)}, nullptr, &n)); + ret.emplace_back(MakeNode("zeros_like", n->attrs.name + "_backward_grad_grad_in", + {n->inputs[1]}, nullptr, &n)); + return ret; + }); + // sign MXNET_OPERATOR_REGISTER_UNARY_WITH_RSP_CSR(sign, cpu, mshadow_op::sign) From 32c934699e6e6a1b3ec2a05258ab84a5b7a6b511 Mon Sep 17 00:00:00 2001 From: kshitij12345 Date: Sat, 29 Jun 2019 12:53:38 +0530 Subject: [PATCH 2/3] add relevant tests --- .../python/unittest/test_higher_order_grad.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tests/python/unittest/test_higher_order_grad.py b/tests/python/unittest/test_higher_order_grad.py index 4f1ea9a6c7b8..a2308b0a02c7 100644 --- a/tests/python/unittest/test_higher_order_grad.py +++ b/tests/python/unittest/test_higher_order_grad.py @@ -106,6 +106,35 @@ def grad_grad_op(x): check_second_order_unary(array, log10, grad_grad_op) +@with_seed() +def test_reciprocal(): + def reciprocal(x): + return nd.reciprocal(x) + + def grad_grad_op(x): + return 2/x**3 + + for dim in range(1, 5): + shape = rand_shape_nd(dim) + array = random_arrays(shape) + check_second_order_unary(array, reciprocal, grad_grad_op) + + +@with_seed() +def test_abs(): + def abs(x): + return nd.abs(x) + + def grad_grad_op(x): + return nd.zeros_like(x) + + for dim in range(1, 5): + shape = rand_shape_nd(dim) + array = random_arrays(shape) + check_second_order_unary(array, abs, grad_grad_op) + + + def check_second_order_unary(x, op, grad_grad_op): x = nd.array(x) grad_grad_x = grad_grad_op(x) From ced9e30e91922bd3f205ed44340e63aa9baa0599 Mon Sep 17 00:00:00 2001 From: kshitij12345 Date: Tue, 2 Jul 2019 22:08:48 +0530 Subject: [PATCH 3/3] address comments * fix extra line in tests. * fix missing space. * fix incorrect comment. --- src/operator/tensor/elemwise_unary_op_basic.cc | 2 +- tests/python/unittest/test_higher_order_grad.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/operator/tensor/elemwise_unary_op_basic.cc b/src/operator/tensor/elemwise_unary_op_basic.cc index 29e45ce396ad..476320c67249 100644 --- a/src/operator/tensor/elemwise_unary_op_basic.cc +++ b/src/operator/tensor/elemwise_unary_op_basic.cc @@ -744,7 +744,7 @@ MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU(_backward_abs, unary_bwd& ograds) { // ograds[0]: dL/dxgrad // inputs[0]: dL/dy - // inputs[1]: y + // inputs[1]: x // f(x) -> abs(x) // f'(x) = 1 if x > 0 else -1 // f''(x) = 0 diff --git a/tests/python/unittest/test_higher_order_grad.py b/tests/python/unittest/test_higher_order_grad.py index a2308b0a02c7..3a255468029d 100644 --- a/tests/python/unittest/test_higher_order_grad.py +++ b/tests/python/unittest/test_higher_order_grad.py @@ -112,7 +112,7 @@ def reciprocal(x): return nd.reciprocal(x) def grad_grad_op(x): - return 2/x**3 + return 2 / x**3 for dim in range(1, 5): shape = rand_shape_nd(dim) @@ -134,7 +134,6 @@ def grad_grad_op(x): check_second_order_unary(array, abs, grad_grad_op) - def check_second_order_unary(x, op, grad_grad_op): x = nd.array(x) grad_grad_x = grad_grad_op(x)