add scatter_nd op and scatter_nd_add op (#19571)
* add scatter_nd op, test=document_preview test=develop * fixed the document, test=document_preview test=develop * modify the notes, test=document_preview test=develop * remove the ShareDataWith, test=developsigmoid_bug
ShenLiang
6 years ago
committed by
Yi Liu
parent
364c44422e
commit
2cd3fa3e9a
@ -0,0 +1,186 @@
|
||||
/* Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
|
||||
|
||||
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
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
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. */
|
||||
|
||||
#include "paddle/fluid/operators/scatter_nd_add_op.h"
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
#include "paddle/fluid/framework/ddim.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace operators {
|
||||
|
||||
class ScatterNdAddOp : public framework::OperatorWithKernel {
|
||||
public:
|
||||
using framework::OperatorWithKernel::OperatorWithKernel;
|
||||
|
||||
void InferShape(framework::InferShapeContext* ctx) const override {
|
||||
PADDLE_ENFORCE_EQ(ctx->HasInput("X"), true,
|
||||
"Input(X) of ScatterNdAddOp should not be null.");
|
||||
PADDLE_ENFORCE_EQ(ctx->HasInput("Index"), true,
|
||||
"Input(Index) of ScatterNdAddOp should not be null.");
|
||||
PADDLE_ENFORCE_EQ(ctx->HasInput("Updates"), true,
|
||||
"Input(Updates) of ScatterNdAddOp should not be null.");
|
||||
PADDLE_ENFORCE_EQ(ctx->HasOutput("Out"), true,
|
||||
"Output(Out) of ScatterNdAddOp should not be null.");
|
||||
|
||||
auto ref_dims = ctx->GetInputDim("X");
|
||||
auto ref_dims_size = ref_dims.size();
|
||||
auto index_dims = ctx->GetInputDim("Index");
|
||||
auto index_dims_size = index_dims.size();
|
||||
auto updates_dims = ctx->GetInputDim("Updates");
|
||||
auto updates_dims_size = updates_dims.size();
|
||||
|
||||
PADDLE_ENFORCE_LE(
|
||||
index_dims[index_dims_size - 1], ref_dims_size,
|
||||
"Input(Index).shape[-1] should be no greater than Input(X).rank");
|
||||
PADDLE_ENFORCE_GE(index_dims_size, 2UL,
|
||||
"The rank of Input(Index) should be greater than 1");
|
||||
|
||||
// update.shape = index.shape[:-1] + output.shape[index.shape[-1]:]
|
||||
std::vector<int64_t> r_updates_dims;
|
||||
for (int64_t i = 0; i < index_dims_size - 1; ++i) {
|
||||
r_updates_dims.emplace_back(index_dims[i]);
|
||||
}
|
||||
for (int64_t i = index_dims[index_dims_size - 1]; i < ref_dims_size; ++i) {
|
||||
r_updates_dims.emplace_back(ref_dims[i]);
|
||||
}
|
||||
|
||||
PADDLE_ENFORCE_EQ(r_updates_dims.size(), updates_dims_size,
|
||||
"Updates has wrong shape");
|
||||
|
||||
for (int64_t i = 0; i < updates_dims_size; ++i) {
|
||||
PADDLE_ENFORCE_EQ(r_updates_dims[i], updates_dims[i],
|
||||
"Updates has wrong shape");
|
||||
}
|
||||
ctx->SetOutputDim("Out", ref_dims);
|
||||
}
|
||||
|
||||
protected:
|
||||
framework::OpKernelType GetExpectedKernelType(
|
||||
const framework::ExecutionContext& ctx) const override {
|
||||
PADDLE_ENFORCE_EQ(ctx.Input<Tensor>("X")->type(),
|
||||
ctx.Input<Tensor>("Updates")->type(),
|
||||
"Ref and Updates must have same type");
|
||||
return framework::OpKernelType(ctx.Input<Tensor>("X")->type(),
|
||||
ctx.device_context());
|
||||
}
|
||||
};
|
||||
|
||||
class ScatterNdAddGradOp : public framework::OperatorWithKernel {
|
||||
public:
|
||||
using framework::OperatorWithKernel::OperatorWithKernel;
|
||||
|
||||
void InferShape(framework::InferShapeContext* ctx) const override {
|
||||
if (ctx->HasOutput(framework::GradVarName("Updates"))) {
|
||||
ctx->SetOutputDim(framework::GradVarName("Updates"),
|
||||
ctx->GetInputDim("Updates"));
|
||||
}
|
||||
if (ctx->HasOutput(framework::GradVarName("X"))) {
|
||||
ctx->SetOutputDim(framework::GradVarName("X"),
|
||||
ctx->GetInputDim(framework::GradVarName("Out")));
|
||||
}
|
||||
}
|
||||
|
||||
protected:
|
||||
framework::OpKernelType GetExpectedKernelType(
|
||||
const framework::ExecutionContext& ctx) const override {
|
||||
return framework::OpKernelType(
|
||||
ctx.Input<Tensor>(framework::GradVarName("Out"))->type(),
|
||||
ctx.device_context());
|
||||
}
|
||||
};
|
||||
|
||||
class ScatterNdAddOpMaker : public framework::OpProtoAndCheckerMaker {
|
||||
public:
|
||||
void Make() override {
|
||||
AddInput("X", "The source input of scatter_nd_add op");
|
||||
AddInput("Index",
|
||||
"The index input of scatter_nd_add op where X will be updated");
|
||||
AddInput("Updates", "The updated value of scatter_nd_add op");
|
||||
AddOutput("Out", "The output of scatter_nd_add op");
|
||||
AddComment(R"DOC(
|
||||
Scatter_nd_add Operator.
|
||||
|
||||
Output is obtained by applying sparse addition to a single value or slice in a Variable.
|
||||
|
||||
Given:
|
||||
* Case 1:
|
||||
ref = [0, 1, 2, 3, 4, 5]
|
||||
index = [[1], [2], [3], [1]]
|
||||
updates = [9, 10, 11, 12]
|
||||
|
||||
we get:
|
||||
|
||||
output = [0, 22, 12, 14, 4, 5]
|
||||
|
||||
* Case 2:
|
||||
ref = [[65, 17], [-14, -25]]
|
||||
index = [[], []]
|
||||
updates = [[[-1, -2], [1, 2]],
|
||||
[[3, 4], [-3, -4]]]
|
||||
ref.shape = (2, 2)
|
||||
index.shape = (2, 0)
|
||||
updates.shape = (2, 2, 2)
|
||||
|
||||
we get:
|
||||
|
||||
output = [[67, 19], [-16, -27]]
|
||||
)DOC");
|
||||
}
|
||||
};
|
||||
|
||||
class ScatterNdAddGradDescMaker : 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_nd_add_grad");
|
||||
op->SetInput("Index", Input("Index"));
|
||||
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(ScatterNdAddGradNoNeedBufferVarsInference,
|
||||
"Updates");
|
||||
|
||||
} // namespace operators
|
||||
} // namespace paddle
|
||||
|
||||
namespace ops = paddle::operators;
|
||||
|
||||
REGISTER_OPERATOR(scatter_nd_add, ops::ScatterNdAddOp, ops::ScatterNdAddOpMaker,
|
||||
ops::ScatterNdAddGradDescMaker);
|
||||
|
||||
REGISTER_OPERATOR(scatter_nd_add_grad, ops::ScatterNdAddGradOp,
|
||||
ops::ScatterNdAddGradNoNeedBufferVarsInference);
|
||||
|
||||
REGISTER_OP_CPU_KERNEL(scatter_nd_add, ops::ScatterNdAddOpKernel<float>,
|
||||
ops::ScatterNdAddOpKernel<double>,
|
||||
ops::ScatterNdAddOpKernel<int64_t>,
|
||||
ops::ScatterNdAddOpKernel<int>,
|
||||
ops::ScatterNdAddOpKernel<uint8_t>);
|
||||
|
||||
REGISTER_OP_CPU_KERNEL(scatter_nd_add_grad,
|
||||
ops::ScatterNdAddGradientOpKernel<float>,
|
||||
ops::ScatterNdAddGradientOpKernel<double>,
|
||||
ops::ScatterNdAddGradientOpKernel<int64_t>,
|
||||
ops::ScatterNdAddGradientOpKernel<int>,
|
||||
ops::ScatterNdAddGradientOpKernel<uint8_t>);
|
||||
@ -0,0 +1,98 @@
|
||||
/* Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
|
||||
|
||||
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
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
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. */
|
||||
|
||||
#include "paddle/fluid/operators/gather.cu.h"
|
||||
#include "paddle/fluid/operators/gather_op.h"
|
||||
#include "paddle/fluid/operators/scatter.cu.h"
|
||||
#include "paddle/fluid/operators/scatter_nd_add_op.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace operators {
|
||||
|
||||
template <typename DeviceContext, typename T>
|
||||
class ScatterNdAddOpCUDAKernel : public framework::OpKernel<T> {
|
||||
public:
|
||||
void Compute(const framework::ExecutionContext &ctx) const override {
|
||||
PADDLE_ENFORCE_EQ(platform::is_gpu_place(ctx.GetPlace()), true,
|
||||
"This kernel only runs on GPU device.");
|
||||
auto *X = ctx.Input<Tensor>("X");
|
||||
auto *Ids = ctx.Input<Tensor>("Index");
|
||||
auto *Updates = ctx.Input<Tensor>("Updates");
|
||||
auto *Out = ctx.Output<Tensor>("Out");
|
||||
|
||||
framework::TensorCopySync(*X, ctx.GetPlace(), Out);
|
||||
const auto &index_type = Ids->type();
|
||||
bool index_type_match = index_type == framework::proto::VarType::INT32 ||
|
||||
index_type == framework::proto::VarType::INT64;
|
||||
PADDLE_ENFORCE_EQ(
|
||||
index_type_match, true,
|
||||
"Index holds the wrong type, it holds %s, but desires to be %s or %s",
|
||||
paddle::framework::DataTypeToString(index_type),
|
||||
paddle::framework::DataTypeToString(framework::proto::VarType::INT32),
|
||||
paddle::framework::DataTypeToString(framework::proto::VarType::INT64));
|
||||
if (index_type == framework::proto::VarType::INT32) {
|
||||
GPUScatterNdAdd<DeviceContext, T, int32_t>(ctx, *Updates, *Ids, Out);
|
||||
} else {
|
||||
GPUScatterNdAdd<DeviceContext, T, int64_t>(ctx, *Updates, *Ids, Out);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <typename DeviceContext, typename T>
|
||||
class ScatterNdAddGradOpCUDAKernel : public framework::OpKernel<T> {
|
||||
public:
|
||||
void Compute(const framework::ExecutionContext &ctx) const override {
|
||||
PADDLE_ENFORCE_EQ(platform::is_gpu_place(ctx.GetPlace()), true,
|
||||
"This kernel only runs on GPU device.");
|
||||
auto *dX = ctx.Output<Tensor>(framework::GradVarName("X"));
|
||||
auto *dUpdates = ctx.Output<Tensor>(framework::GradVarName("Updates"));
|
||||
auto *Ids = ctx.Input<Tensor>("Index");
|
||||
auto *dOut = ctx.Input<Tensor>(framework::GradVarName("Out"));
|
||||
if (dX) {
|
||||
// In place gradient: dX = dO
|
||||
framework::TensorCopy(*dOut, ctx.GetPlace(), dX);
|
||||
}
|
||||
if (dUpdates) {
|
||||
dUpdates->mutable_data<T>(ctx.GetPlace());
|
||||
// Gradient by Gather
|
||||
const auto &index_type = Ids->type();
|
||||
if (index_type == framework::proto::VarType::INT32) {
|
||||
GPUGatherNd<DeviceContext, T, int32_t>(ctx, *dOut, *Ids, dUpdates);
|
||||
} else {
|
||||
GPUGatherNd<DeviceContext, T, int64_t>(ctx, *dOut, *Ids, dUpdates);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace operators
|
||||
} // namespace paddle
|
||||
|
||||
namespace ops = paddle::operators;
|
||||
using CUDA = paddle::platform::CUDADeviceContext;
|
||||
namespace plat = paddle::platform;
|
||||
|
||||
REGISTER_OP_CUDA_KERNEL(scatter_nd_add,
|
||||
ops::ScatterNdAddOpCUDAKernel<CUDA, float>,
|
||||
ops::ScatterNdAddOpCUDAKernel<CUDA, double>,
|
||||
ops::ScatterNdAddOpCUDAKernel<CUDA, int64_t>,
|
||||
ops::ScatterNdAddOpCUDAKernel<CUDA, int>,
|
||||
ops::ScatterNdAddOpCUDAKernel<CUDA, plat::float16>);
|
||||
|
||||
REGISTER_OP_CUDA_KERNEL(scatter_nd_add_grad,
|
||||
ops::ScatterNdAddGradOpCUDAKernel<CUDA, float>,
|
||||
ops::ScatterNdAddGradOpCUDAKernel<CUDA, double>,
|
||||
ops::ScatterNdAddGradOpCUDAKernel<CUDA, int64_t>,
|
||||
ops::ScatterNdAddGradOpCUDAKernel<CUDA, int>,
|
||||
ops::ScatterNdAddGradOpCUDAKernel<CUDA, plat::float16>);
|
||||
@ -0,0 +1,86 @@
|
||||
/* Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
|
||||
|
||||
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
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
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. */
|
||||
|
||||
#pragma once
|
||||
#include "paddle/fluid/framework/eigen.h"
|
||||
#include "paddle/fluid/framework/op_registry.h"
|
||||
#include "paddle/fluid/operators/gather.h"
|
||||
#include "paddle/fluid/operators/scatter.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace operators {
|
||||
|
||||
using Tensor = framework::Tensor;
|
||||
|
||||
template <typename T>
|
||||
class ScatterNdAddOpKernel : public framework::OpKernel<T> {
|
||||
public:
|
||||
void Compute(const framework::ExecutionContext &ctx) const override {
|
||||
PADDLE_ENFORCE_EQ(platform::is_cpu_place(ctx.GetPlace()), true,
|
||||
"This kernel only runs on CPU.");
|
||||
auto *X = ctx.Input<Tensor>("X");
|
||||
auto *Ids = ctx.Input<Tensor>("Index");
|
||||
auto *Updates = ctx.Input<Tensor>("Updates");
|
||||
auto *Out = ctx.Output<Tensor>("Out");
|
||||
|
||||
// In place output: Out = X
|
||||
framework::TensorCopySync(*X, ctx.GetPlace(), Out);
|
||||
const auto &index_type = Ids->type();
|
||||
bool index_type_match = index_type == framework::proto::VarType::INT32 ||
|
||||
index_type == framework::proto::VarType::INT64;
|
||||
PADDLE_ENFORCE_EQ(
|
||||
index_type_match, true,
|
||||
"Index holds the wrong type, it holds %s, but desires to be %s or %s",
|
||||
paddle::framework::DataTypeToString(index_type),
|
||||
paddle::framework::DataTypeToString(framework::proto::VarType::INT32),
|
||||
paddle::framework::DataTypeToString(framework::proto::VarType::INT64));
|
||||
|
||||
if (index_type == framework::proto::VarType::INT32) {
|
||||
ScatterNdAdd<T, int32_t>(ctx, *Updates, *Ids, Out);
|
||||
} else {
|
||||
ScatterNdAdd<T, int64_t>(ctx, *Updates, *Ids, Out);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
class ScatterNdAddGradientOpKernel : public framework::OpKernel<T> {
|
||||
public:
|
||||
void Compute(const framework::ExecutionContext &ctx) const override {
|
||||
PADDLE_ENFORCE_EQ(platform::is_cpu_place(ctx.GetPlace()), true,
|
||||
"This kernel only runs on CPU.");
|
||||
auto *dX = ctx.Output<Tensor>(framework::GradVarName("X"));
|
||||
auto *dUpdates = ctx.Output<Tensor>(framework::GradVarName("Updates"));
|
||||
auto *Ids = ctx.Input<Tensor>("Index");
|
||||
auto *dOut = ctx.Input<Tensor>(framework::GradVarName("Out"));
|
||||
|
||||
if (dX) {
|
||||
// In place gradient: dX = dO
|
||||
framework::TensorCopySync(*dOut, ctx.GetPlace(), dX);
|
||||
}
|
||||
if (dUpdates) {
|
||||
dUpdates->mutable_data<T>(ctx.GetPlace());
|
||||
// Gradient by Gather: dUpdates = dO[Ids]
|
||||
const auto &index_type = Ids->type();
|
||||
if (index_type == framework::proto::VarType::INT32) {
|
||||
CPUGatherNd<T, int32_t>(ctx.device_context(), *dOut, *Ids, dUpdates);
|
||||
} else {
|
||||
CPUGatherNd<T, int64_t>(ctx.device_context(), *dOut, *Ids, dUpdates);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace operators
|
||||
} // namespace paddle
|
||||
File diff suppressed because it is too large
Load Diff
Loading…
Reference in new issue