add higher order gradient support for log, log10, log2

apache · May 23, 2019 · e3ce4ab · e3ce4ab
1 parent d2e397a
commit e3ce4ab
Showing 1 changed file with 62 additions and 3 deletions.
diff --git a/src/operator/tensor/elemwise_unary_op_basic.cc b/src/operator/tensor/elemwise_unary_op_basic.cc
@@ -1069,13 +1069,72 @@ The storage type of ``log2`` output is always dense
 .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_log2"});
 
 MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU_DR(_backward_log,
-                                                  unary_bwd<mshadow_op::log_grad>);
+                                                  unary_bwd<mshadow_op::log_grad>)
+.set_attr<nnvm::FGradient>("FGradient",
+  [](const nnvm::NodePtr& n, const std::vector<nnvm::NodeEntry>& ograds) {
+    // For g(x) -> g = log
+    // g''(x) = -1 * (g'(x) * g'(x))
+    auto gx = nnvm::NodeEntry{n, 0, 0};
+    auto ggx_mid = MakeNode("elemwise_mul", n->attrs.name + "_backward_mid_grad_grad",
+                            {gx, gx}, nullptr, &n);
+    auto ggx = MakeNode("negative", n->attrs.name + "_backward_grad_grad",
+                        {nnvm::NodeEntry{ggx_mid, 0, 0}}, nullptr, &n);
+
+    std::vector<nnvm::NodeEntry> ret;
+
+    ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad",
+                             {ograds[0], gx}, nullptr, &n));
+    ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad_inp",
+                             {ograds[0], nnvm::NodeEntry{ggx, 0, 0}}, nullptr, &n));
+    return ret;
+  });
 
 MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU_DR(_backward_log10,
-                                                  unary_bwd<mshadow_op::log10_grad>);
+                                                  unary_bwd<mshadow_op::log10_grad>)
+.set_attr<nnvm::FGradient>("FGradient",
+  [](const nnvm::NodePtr& n, const std::vector<nnvm::NodeEntry>& ograds) {
+    // For g(x) -> g = log10
+    // g'(x) = 1 / (log(10) * x)
+    // g''(x) = -1 * (g'(x) * 1/x)
+    auto gx = nnvm::NodeEntry{n, 0, 0};
+    auto g_lx = MakeNode("reciprocal", n->attrs.name + "_backward_log_grad",
+                            {n->inputs[1]}, nullptr, &n);
+    auto ggx_mid = MakeNode("elemwise_mul", n->attrs.name + "_backward_mid_grad_grad",
+                            {gx, nnvm::NodeEntry{g_lx, 0, 0}}, nullptr, &n);
+    auto ggx = MakeNode("negative", n->attrs.name + "_backward_grad_grad",
+                        {nnvm::NodeEntry{ggx_mid, 0, 0}}, nullptr, &n);
+
+    std::vector<nnvm::NodeEntry> ret;
+
+    ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad",
+                             {ograds[0], gx}, nullptr, &n));
+    ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad_inp",
+                             {ograds[0], nnvm::NodeEntry{ggx, 0, 0}}, nullptr, &n));
+    return ret;
+  });
 
 MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU_DR(_backward_log2,
-                                                  unary_bwd<mshadow_op::log2_grad>);
+                                                  unary_bwd<mshadow_op::log2_grad>)
+.set_attr<nnvm::FGradient>("FGradient",
+  [](const nnvm::NodePtr& n, const std::vector<nnvm::NodeEntry>& ograds) {
+    // For g(x) -> g = log2
+    // g'(x) = 1 / (log(2) * x)
+    // g''(x) = -1 * (g'(x) * 1/x)
+    auto gx = nnvm::NodeEntry{n, 0, 0};
+    auto g_lx = MakeNode("reciprocal", n->attrs.name + "_backward_log_grad",
+                            {n->inputs[1]}, nullptr, &n);
+    auto ggx_mid = MakeNode("elemwise_mul", n->attrs.name + "_backward_mid_grad_grad",
+                            {gx, nnvm::NodeEntry{g_lx, 0, 0}}, nullptr, &n);
+    auto ggx = MakeNode("negative", n->attrs.name + "_backward_grad_grad",
+                        {nnvm::NodeEntry{ggx_mid, 0, 0}}, nullptr, &n);
+
+    std::vector<nnvm::NodeEntry> ret;
+    ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad",
+                             {ograds[0], gx}, nullptr, &n));
+    ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad_inp",
+                             {ograds[0], nnvm::NodeEntry{ggx, 0, 0}}, nullptr, &n));
+    return ret;
+  });
 
 // log1p
 MXNET_OPERATOR_REGISTER_UNARY_WITH_RSP_CSR(log1p, cpu, mshadow_op::log1p)