Add Gradient Operator for mean

paddle/operators/mean_op.cc

@@ -33,13 +33,23 @@ class MeanOpMaker : public OpProtoAndCheckerMaker {
  MeanOpMaker(OpProto *proto, OpAttrChecker *op_checker)
  : OpProtoAndCheckerMaker(proto, op_checker) {
  AddInput("X", "The input of mean op");
- AddOutput("Out", "The output of mean op");
+ AddOutput("Out", "The output of mean op").IgnoreGradient();
  AddComment("Mean Operator");
  }
 };
 
+class MeanGradOp : public OperatorWithKernel {
+protected:
+ void InferShape(const InferShapeContext &ctx) const override {
+ ctx.Output<Tensor>("X" + GRAD_VAR_SUFFIX())
+ ->Resize(ctx.Input<Tensor>("X")->dims());
+ }
+};
+
 } // namespace operators
 } // namespace paddle
 
 REGISTER_OP(mean, ops::MeanOp, ops::MeanOpMaker);
 REGISTER_OP_CPU_KERNEL(mean, ops::MeanKernel<ops::CPUPlace, float>);
+REGISTER_GRADIENT_OP(mean, mean_grad, ops::MeanGradOp);
+REGISTER_OP_CPU_KERNEL(mean_grad, ops::MeanGradKernel<ops::CPUPlace, float>);

paddle/operators/mean_op.cu

@@ -3,3 +3,4 @@
 #include "paddle/operators/mean_op.h"
 
 REGISTER_OP_GPU_KERNEL(mean, ops::MeanKernel<ops::GPUPlace, float>);
+REGISTER_OP_GPU_KERNEL(mean_grad, ops::MeanGradKernel<ops::GPUPlace, float>);

paddle/operators/mean_op.h

@@ -32,5 +32,22 @@ class MeanKernel : public OpKernel {
  }
 };
 
+template <typename Place, typename T>
+class MeanGradKernel : public OpKernel {
+public:
+ void Compute(const ExecutionContext& context) const override {
+ auto OG = context.Input<Tensor>("Out" + OperatorBase::GRAD_VAR_SUFFIX());
+ PADDLE_ENFORCE(framework::product(OG->dims()) == 1,
+ "Mean Gradient should be scalar");
+ auto IG = context.Output<Tensor>("X" + OperatorBase::GRAD_VAR_SUFFIX());
+ IG->mutable_data<T>(context.GetPlace());
+
+ T ig_size = (T)framework::product(IG->dims());
+
+ EigenVector<T>::Flatten(*IG).device(*(context.GetEigenDevice<Place>())) =
+ EigenScalar<T>::From(*OG) / ig_size;
+ }
+};
+
 } // namespace operators
 } // namespace paddle

paddle/operators/type_alias.h

@@ -51,6 +51,7 @@ using CPUPlace = platform::CPUPlace;
 using GPUPlace = platform::GPUPlace;
 using NetOp = framework::NetOp;
 using OpRegistry = framework::OpRegistry;
+using OperatorBase = framework::OperatorBase;
 } // namespace operators
 } // namespace paddle