scatter_op.cc 5.0 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
Z
zchen0211 已提交
2

Z
zchen0211 已提交
3 4 5
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
Z
zchen0211 已提交
6

Z
zchen0211 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
Z
zchen0211 已提交
8

Z
zchen0211 已提交
9 10 11 12 13
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */
Z
zchen0211 已提交
14

Y
Yi Wang 已提交
15
#include "paddle/fluid/operators/scatter_op.h"
S
sneaxiy 已提交
16
#include <memory>
Y
Yi Wang 已提交
17
#include "paddle/fluid/framework/ddim.h"
Z
zchen0211 已提交
18 19 20 21 22 23 24 25

namespace paddle {
namespace operators {

class ScatterOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

26
  void InferShape(framework::InferShapeContext* ctx) const override {
D
dzhwinter 已提交
27 28 29 30
    PADDLE_ENFORCE(ctx->HasInput("X"),
                   "Input(X) of ScatterOp should not be null.");
    PADDLE_ENFORCE(ctx->HasInput("Ids"),
                   "Input(Ids) of ScatterOp should not be null.");
Q
Qiao Longfei 已提交
31 32 33 34
    PADDLE_ENFORCE(ctx->HasInput("Updates"),
                   "Input(Updates) of ScatterOp should not be null.");
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   "Output(Out) of ScatterOp should not be null.");
35

Q
Qiao Longfei 已提交
36
    auto updates_dims = ctx->GetInputDim("Updates");
D
dzhwinter 已提交
37 38 39
    auto ref_dims = ctx->GetInputDim("X");
    PADDLE_ENFORCE_EQ(ctx->GetInputDim("Ids").size(), 1,
                      "Update Ids should be 1-D.");
Q
Qiao Longfei 已提交
40
    PADDLE_ENFORCE_EQ(ref_dims.size(), updates_dims.size(),
D
dzhwinter 已提交
41
                      "Xerence and Updates should have the same shape size");
Q
Qiao Longfei 已提交
42
    PADDLE_ENFORCE_EQ(ctx->GetInputDim("Updates")[0],
D
dzhwinter 已提交
43 44
                      ctx->GetInputDim("Ids")[0],
                      "Updates and Ids should have same batch-size.");
Q
Qiao Longfei 已提交
45
    ctx->SetOutputDim("Out", ref_dims);
Z
zchen0211 已提交
46
  }
Y
Yu Yang 已提交
47

48
 protected:
49
  framework::OpKernelType GetExpectedKernelType(
Y
Yu Yang 已提交
50
      const framework::ExecutionContext& ctx) const override {
Y
Yu Yang 已提交
51 52
    return framework::OpKernelType(ctx.Input<Tensor>("X")->type(),
                                   ctx.device_context());
Y
Yu Yang 已提交
53
  }
Z
zchen0211 已提交
54 55 56 57 58 59
};

class ScatterGradOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

60
  void InferShape(framework::InferShapeContext* ctx) const override {
Q
Qiao Longfei 已提交
61 62
    ctx->SetOutputDim(framework::GradVarName("Updates"),
                      ctx->GetInputDim("Updates"));
S
sneaxiy 已提交
63 64
    ctx->SetOutputDim(framework::GradVarName("X"),
                      ctx->GetInputDim(framework::GradVarName("Out")));
Z
zchen0211 已提交
65
  }
Y
Yu Yang 已提交
66

67
 protected:
68
  framework::OpKernelType GetExpectedKernelType(
Y
Yu Yang 已提交
69
      const framework::ExecutionContext& ctx) const override {
S
sneaxiy 已提交
70 71 72
    return framework::OpKernelType(
        ctx.Input<Tensor>(framework::GradVarName("Out"))->type(),
        ctx.device_context());
Y
Yu Yang 已提交
73
  }
Z
zchen0211 已提交
74 75 76 77
};

class ScatterOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
78
  void Make() override {
D
dzhwinter 已提交
79 80
    AddInput("X", "The source input of scatter op");
    AddInput("Ids", "The index input of scatter op where X will be updated");
81 82
    AddInput("Updates", "The updated value of scatter op");
    AddOutput("Out", "The output of scatter op");
83 84 85 86 87 88 89 90
    AddAttr<bool>("overwrite",
                  "(bool, defalut: True) "
                  "The mode that updating the output when has same index,"
                  "If True, use the overwrite mode to update the output"
                  "of the same index, if False, use the accumulate mode to"
                  "update the output of the same index,Default value is True."
                  "You can set overwrite=False to implement scatter_add.")
        .SetDefault(true);
Z
zchen0211 已提交
91
    AddComment(R"DOC(
K
kexinzhao 已提交
92
Scatter Operator.
Z
zchen0211 已提交
93

K
kexinzhao 已提交
94 95 96
This operator obtains output by updating the input on selected indices on the first axis:

$$
D
dzhwinter 已提交
97
Out = X \\
98
Out[Ids] = Updates
K
kexinzhao 已提交
99 100
$$

Z
zchen0211 已提交
101 102 103
)DOC");
  }
};
Q
Qiao Longfei 已提交
104

S
sneaxiy 已提交
105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125
class ScatterGradDescMaker : public framework::SingleGradOpDescMaker {
 public:
  using framework::SingleGradOpDescMaker::SingleGradOpDescMaker;

 protected:
  std::unique_ptr<framework::OpDesc> Apply() const override {
    std::unique_ptr<framework::OpDesc> op(new framework::OpDesc());
    op->SetType("scatter_grad");
    op->SetInput("Ids", Input("Ids"));
    op->SetInput("Updates", Input("Updates"));
    op->SetInput(framework::GradVarName("Out"), OutputGrad("Out"));
    op->SetOutput(framework::GradVarName("X"), InputGrad("X"));
    op->SetOutput(framework::GradVarName("Updates"), InputGrad("Updates"));
    op->SetAttrMap(Attrs());
    return op;
  }
};

DECLARE_NO_NEED_BUFFER_VARS_INFERENCE(ScatterGradNoNeedBufferVarsInference,
                                      "Updates");

Z
zchen0211 已提交
126 127 128 129
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
Y
Yang Yang 已提交
130
REGISTER_OPERATOR(scatter, ops::ScatterOp, ops::ScatterOpMaker,
S
sneaxiy 已提交
131 132 133
                  ops::ScatterGradDescMaker);
REGISTER_OPERATOR(scatter_grad, ops::ScatterGradOp,
                  ops::ScatterGradNoNeedBufferVarsInference);
Z
zchen0211 已提交
134 135
REGISTER_OP_CPU_KERNEL(scatter, ops::ScatterOpKernel<float>);
REGISTER_OP_CPU_KERNEL(scatter_grad, ops::ScatterGradientOpKernel<float>);