Revert D6893040: Implementing Pow operator (this merges existing pow …

…with a scalar and new pow with a tensor exponent). Summary: This reverts commit 30f614beea6f859fee25ce4f85573142885dde45 bypass-lint An infra SEV is better than not reverting this diff. If you copy this password, see you in SEV Review! cause_a_sev_many_files Differential Revision: D6893040 Original commit changeset: 30f614beea6f fbshipit-source-id: 5e98a24699088283f864efe31234874bdacbe3c3
duc0 · Feb 14, 2018 · 52fa742 · 52fa742
1 parent c746357
commit 52fa742
Show file tree

Hide file tree

Showing 9 changed files with 81 additions and 580 deletions.
diff --git a/caffe2/operators/elementwise_op.cu b/caffe2/operators/elementwise_op.cu
@@ -113,16 +113,6 @@ CUDA_FUNCTOR(Or, CUDA_OR, BoolTypes, FixedType<bool>);
 CUDA_FUNCTOR(Xor, CUDA_XOR, BoolTypes, FixedType<bool>);
 #undef CUDA_XOR
 
-// pow, log and other math functions are defined in CUDA math library
-// in header file math.h
-#define CUDA_POW(x, y) (pow(x, y))
-CUDA_FUNCTOR(
-    Pow,
-    CUDA_POW,
-    TensorTypes<float> /*NumericTypes*/,
-    SameTypeAsInput);
-#undef CUDA_POW
-
 __global__ void NotKernel(const int n, const bool* x, bool* y) {
   CUDA_1D_KERNEL_LOOP(i, n) {
     y[i] = !x[i];

diff --git a/caffe2/operators/math_ops.cc b/caffe2/operators/math_ops.cc
@@ -82,4 +82,65 @@ OPERATOR_SCHEMA(Sign)
     .IdenticalTypeAndShape();
 SHOULD_NOT_DO_GRADIENT(Sign);
 
+REGISTER_CPU_OPERATOR(
+    Pow,
+    UnaryElementwiseWithArgsOp<TensorTypes<float>, CPUContext, PowFunctor>);
+
+OPERATOR_SCHEMA(Pow)
+    .NumInputs(1)
+    .NumOutputs(1)
+    .Arg("exponent", "The exponent of the power function.")
+    .AllowInplace({{0, 0}})
+    .IdenticalTypeAndShape()
+    .SetDoc(R"DOC(
+Pow takes input data (Tensor<T>) and an argument exponent, and
+produces one output data (Tensor<T>) where the function `f(x) = x^exponent`,
+is applied to the data tensor elementwise.
+)DOC")
+    .Input(0, "X", "Input tensor of any shape")
+    .Output(0, "Y", "Output tensor (same size as X)");
+
+class GetPowGradient : public GradientMakerBase {
+  using GradientMakerBase::GradientMakerBase;
+  vector<OperatorDef> GetGradientDefs() override {
+    ArgumentHelper arg_helper(def_);
+    float exponent = arg_helper.GetSingleArgument<float>("exponent", 0.0);
+    Argument scale_arg;
+    scale_arg.set_name("scale");
+    scale_arg.set_f(exponent);
+    Argument pow_arg;
+    pow_arg.set_name("exponent");
+    if (I(0) != O(0)) {
+      pow_arg.set_f(exponent - 1);
+    } else {
+      LOG(WARNING) << "In-place Pow gradient, possible loss of precision";
+      constexpr float kEps = 1e-12f;
+      CAFFE_ENFORCE(std::fabs(exponent) > kEps);
+      pow_arg.set_f((exponent - 1) / exponent);
+    }
+    return vector<OperatorDef>{CreateOperatorDef(
+                                   "Pow",
+                                   "",
+                                   std::vector<string>{I(0)},
+                                   std::vector<string>{GI(0)},
+                                   std::vector<Argument>{pow_arg}),
+                               CreateOperatorDef(
+                                   "Mul",
+                                   "",
+                                   std::vector<string>{GI(0), GO(0)},
+                                   std::vector<string>{GI(0)}),
+                               CreateOperatorDef(
+                                   "Scale",
+                                   "",
+                                   std::vector<string>{GI(0)},
+                                   std::vector<string>{GI(0)},
+                                   std::vector<Argument>{scale_arg})};
+  }
+  virtual bool CopyArguments() const override {
+    return false;
+  }
+};
+
+REGISTER_GRADIENT(Pow, GetPowGradient);
+
 } // namespace caffe2
diff --git a/caffe2/operators/math_ops.cu b/caffe2/operators/math_ops.cu
@@ -52,4 +52,7 @@ REGISTER_CUDA_OPERATOR(
 REGISTER_CUDA_OPERATOR(
     Sign,
     UnaryElementwiseOp<TensorTypes<float>, CUDAContext, SignCUDAFunctor>);
+REGISTER_CUDA_OPERATOR(
+    Pow,
+    UnaryElementwiseWithArgsOp<TensorTypes<float>, CUDAContext, PowFunctor>);
 }
diff --git a/caffe2/operators/math_ops.h b/caffe2/operators/math_ops.h
@@ -25,4 +25,21 @@
 #include "caffe2/operators/elementwise_op.h"
 #include "caffe2/utils/math.h"
 
+namespace caffe2 {
+
+struct PowFunctor {
+  explicit PowFunctor(OperatorBase& op) {
+    exponent_ = op.GetSingleArgument<float>("exponent", 0);
+  }
+
+  template <typename T, class Context>
+  inline void
+  operator()(const int n, const T* x, T* y, Context* device_context) {
+    math::Powx<float, Context>(n, x, exponent_, y, device_context);
+  }
+
+  float exponent_;
+};
+}
+
 #endif