|
|
/* Copyright (c) 2018 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/sequence_ops/sequence_scatter_op.h"
|
|
|
#include <memory>
|
|
|
#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;
|
|
|
using LoDTensor = framework::LoDTensor;
|
|
|
|
|
|
class SequenceScatterOpMaker : public framework::OpProtoAndCheckerMaker {
|
|
|
public:
|
|
|
void Make() override {
|
|
|
AddInput("X", "(Tensor) The source input of sequence scatter op");
|
|
|
AddInput("Ids",
|
|
|
"(LoDTensor) The index input of sequence scatter op where X"
|
|
|
" will be updated, must be a LoDTensor");
|
|
|
AddInput("Updates",
|
|
|
"(LoDTensor) The values to scatter to the input tensor "
|
|
|
"X, must be a LoDTensor with the same LoD information as Ids");
|
|
|
AddOutput("Out",
|
|
|
"(Tensor) The output tensor of sequence scatter op, which "
|
|
|
"has the same dims as X");
|
|
|
AddComment(R"DOC(
|
|
|
Sequence Scatter Operator.
|
|
|
|
|
|
This operator scatters the Updates tensor to the input X. It uses the LoD
|
|
|
information of Ids to select the rows to update, and use the values in Ids as
|
|
|
the columns to update in each row of X.
|
|
|
|
|
|
Following are cases to better explain how this works:
|
|
|
|
|
|
Example 1:
|
|
|
Given an all-ones Tensor input(X)
|
|
|
X.data = [[1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
|
|
|
[1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
|
|
|
[1.0, 1.0, 1.0, 1.0, 1.0, 1.0]]
|
|
|
X.dims = [3, 6]
|
|
|
a LoDTensor input(Ids)
|
|
|
Ids.data = [[0], [1], [2], [5], [4], [3], [2], [1], [3], [2], [5], [4]]
|
|
|
Ids.lod = [[0, 3, 8, 12]]
|
|
|
and a Tensor input(Updates)
|
|
|
Updates.data = [[0.3], [0.3], [0.4], [0.1], [0.2], [0.3], [0.4], [0.0], [0.2], [0.3], [0.1], [0.4]]
|
|
|
Updates.lod = [[ 0, 3, 8, 12]]
|
|
|
then we get an output Tensor
|
|
|
Out.data = [[1.3, 1.3, 1.4, 1.0, 1.0, 1.0],
|
|
|
[1.0, 1.0, 1.4, 1.3, 1.2, 1.1],
|
|
|
[1.0, 1.0, 1.3, 1.2, 1.4, 1.1]]
|
|
|
Out.dims = X.dims = [3, 6]
|
|
|
)DOC");
|
|
|
}
|
|
|
};
|
|
|
|
|
|
class SequenceScatterOp : public framework::OperatorWithKernel {
|
|
|
public:
|
|
|
using framework::OperatorWithKernel::OperatorWithKernel;
|
|
|
|
|
|
void InferShape(framework::InferShapeContext* ctx) const override {
|
|
|
// Enforce has inputs and outputs
|
|
|
OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "SequenceScatter");
|
|
|
OP_INOUT_CHECK(ctx->HasInput("Ids"), "Input", "Ids", "SequenceScatter");
|
|
|
OP_INOUT_CHECK(ctx->HasInput("Updates"), "Input", "Updates",
|
|
|
"SequenceScatter");
|
|
|
OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "SequenceScatter");
|
|
|
|
|
|
// Set output dim the same as input
|
|
|
auto ref_dims = ctx->GetInputDim("X");
|
|
|
ctx->SetOutputDim("Out", ref_dims);
|
|
|
|
|
|
// Enforce the Updates and Ids are the same shape
|
|
|
auto updates_dim = ctx->GetInputDim("Updates");
|
|
|
auto ids_dim = ctx->GetInputDim("Ids");
|
|
|
PADDLE_ENFORCE_EQ(
|
|
|
updates_dim[0], ids_dim[0],
|
|
|
platform::errors::InvalidArgument(
|
|
|
"The shape of SequenceScatter operator's input Updates and Ids do "
|
|
|
"not match, receive Updates's shape is [%s], Ids's shape is [%s].",
|
|
|
updates_dim, ids_dim));
|
|
|
|
|
|
// Enforce LoD of ids and updates be the same
|
|
|
if (ctx->IsRuntime()) {
|
|
|
framework::Variable* ids_var =
|
|
|
BOOST_GET(framework::Variable*, ctx->GetInputVarPtrs("Ids")[0]);
|
|
|
framework::Variable* updates_var =
|
|
|
BOOST_GET(framework::Variable*, ctx->GetInputVarPtrs("Updates")[0]);
|
|
|
|
|
|
auto& ids_lod = ids_var->Get<LoDTensor>().lod();
|
|
|
auto& updates_lod = updates_var->Get<LoDTensor>().lod();
|
|
|
PADDLE_ENFORCE_EQ(
|
|
|
ids_lod.size(), 1,
|
|
|
platform::errors::InvalidArgument(
|
|
|
"The SequenceScatter operator’s Input Ids holds wrong LoD "
|
|
|
"information. Currently SequenceScatter operator can only deal "
|
|
|
"with one level LoD for input Ids, but received LoD level is %d.",
|
|
|
ids_lod.size()));
|
|
|
PADDLE_ENFORCE_EQ(
|
|
|
updates_lod.size(), 1,
|
|
|
platform::errors::InvalidArgument(
|
|
|
"The SequenceScatter operator’s Input Updates holds wrong LoD "
|
|
|
"information. Currently SequenceScatter operator can only deal "
|
|
|
"with one level LoD for input Updates, but received LoD level is "
|
|
|
"%d.",
|
|
|
ids_lod.size()));
|
|
|
}
|
|
|
}
|
|
|
|
|
|
protected:
|
|
|
framework::OpKernelType GetExpectedKernelType(
|
|
|
const framework::ExecutionContext& ctx) const override {
|
|
|
return framework::OpKernelType(
|
|
|
OperatorWithKernel::IndicateVarDataType(ctx, "X"),
|
|
|
platform::CPUPlace());
|
|
|
}
|
|
|
};
|
|
|
|
|
|
class SequenceScatterGradOp : public framework::OperatorWithKernel {
|
|
|
public:
|
|
|
using framework::OperatorWithKernel::OperatorWithKernel;
|
|
|
|
|
|
void InferShape(framework::InferShapeContext* ctx) const override {
|
|
|
ctx->SetOutputDim(framework::GradVarName("Updates"),
|
|
|
ctx->GetInputDim("Updates"));
|
|
|
ctx->SetOutputDim(framework::GradVarName("X"),
|
|
|
ctx->GetInputDim(framework::GradVarName("Out")));
|
|
|
}
|
|
|
|
|
|
protected:
|
|
|
framework::OpKernelType GetExpectedKernelType(
|
|
|
const framework::ExecutionContext& ctx) const override {
|
|
|
return framework::OpKernelType(OperatorWithKernel::IndicateVarDataType(
|
|
|
ctx, framework::GradVarName("Out")),
|
|
|
platform::CPUPlace());
|
|
|
}
|
|
|
};
|
|
|
|
|
|
template <typename T>
|
|
|
class SequenceScatterGradMaker : public framework::SingleGradOpMaker<T> {
|
|
|
public:
|
|
|
using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
|
|
|
|
|
|
protected:
|
|
|
void Apply(GradOpPtr<T> op) const override {
|
|
|
op->SetType("sequence_scatter_grad");
|
|
|
op->SetInput("Ids", this->Input("Ids"));
|
|
|
op->SetInput("Updates", this->Input("Updates"));
|
|
|
op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
|
|
|
op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
|
|
|
op->SetOutput(framework::GradVarName("Updates"),
|
|
|
this->InputGrad("Updates"));
|
|
|
op->SetAttrMap(this->Attrs());
|
|
|
}
|
|
|
};
|
|
|
|
|
|
DECLARE_NO_NEED_BUFFER_VARS_INFERER(SequenceScatterGradNoNeedBufferVarsInferer,
|
|
|
"Updates");
|
|
|
|
|
|
} // namespace operators
|
|
|
} // namespace paddle
|
|
|
|
|
|
namespace ops = paddle::operators;
|
|
|
REGISTER_OPERATOR(sequence_scatter, ops::SequenceScatterOp,
|
|
|
ops::SequenceScatterOpMaker,
|
|
|
ops::SequenceScatterGradMaker<paddle::framework::OpDesc>,
|
|
|
ops::SequenceScatterGradMaker<paddle::imperative::OpBase>);
|
|
|
REGISTER_OPERATOR(sequence_scatter_grad, ops::SequenceScatterGradOp,
|
|
|
ops::SequenceScatterGradNoNeedBufferVarsInferer);
|
|
|
REGISTER_OP_CPU_KERNEL(sequence_scatter, ops::SequenceScatterOpKernel<float>,
|
|
|
ops::SequenceScatterOpKernel<double>,
|
|
|
ops::SequenceScatterOpKernel<int>,
|
|
|
ops::SequenceScatterOpKernel<int64_t>);
|
|
|
REGISTER_OP_CPU_KERNEL(sequence_scatter_grad,
|
|
|
ops::SequenceScatterGradientOpKernel<float>,
|
|
|
ops::SequenceScatterGradientOpKernel<double>,
|
|
|
ops::SequenceScatterGradientOpKernel<int>,
|
|
|
ops::SequenceScatterGradientOpKernel<int64_t>);
|
