Skip to content

Commit 7c1755d

Browse files
reyoungreyoung
authored
Assign Operator. (#5531)
* Assign Operator. Out=X, when type in [LoDTensor/SelectedRows/LoDTensorArray] * Follow comments
1 parent 291f6b4 commit 7c1755d

File tree

3 files changed

+181
-0
lines changed

3 files changed

+181
-0
lines changed

‎paddle/framework/var_type.h

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,10 +27,32 @@ inline VarDesc::VarType ToVarType(std::type_index type) {
2727
return VarDesc_VarType_LOD_RANK_TABLE;
2828
} else if (type.hash_code() == typeid(LoDTensorArray).hash_code()) {
2929
return VarDesc_VarType_LOD_TENSOR_ARRAY;
30+
} else if (type.hash_code() == typeid(SelectedRows).hash_code()) {
31+
return VarDesc_VarType_SELECTED_ROWS;
3032
} else {
3133
PADDLE_THROW("ToVarType:Unsupported type %s", type.name());
3234
}
3335
}
3436

37+
template <typename Visitor>
38+
inline void VisitVarType(const Variable& var, Visitor visitor) {
39+
switch (ToVarType(var.Type())) {
40+
case VarDesc_VarType_LOD_TENSOR:
41+
visitor(var.Get<framework::LoDTensor>());
42+
return;
43+
case VarDesc_VarType_LOD_RANK_TABLE:
44+
visitor(var.Get<LoDRankTable>());
45+
return;
46+
case VarDesc_VarType_LOD_TENSOR_ARRAY:
47+
visitor(var.Get<LoDTensorArray>());
48+
return;
49+
case VarDesc_VarType_SELECTED_ROWS:
50+
visitor(var.Get<SelectedRows>());
51+
return;
52+
default:
53+
PADDLE_THROW("Not supported visit type, %d", ToVarType(var.Type()));
54+
}
55+
}
56+
3557
} // namespace framework
3658
} // namespace paddle

‎paddle/operators/assign_op.cc

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License. */
14+
15+
#include "paddle/framework/data_type.h"
16+
#include "paddle/framework/op_registry.h"
17+
#include "paddle/framework/var_type.h"
18+
19+
namespace paddle {
20+
namespace operators {
21+
class AssignFunctor {
22+
public:
23+
AssignFunctor(framework::Variable *out,
24+
const platform::DeviceContext &dev_ctx)
25+
: out_(out), dev_ctx_(dev_ctx) {}
26+
27+
void operator()(const framework::LoDTensor &lod_tensor) const {
28+
auto &out_tensor = *out_->GetMutable<framework::LoDTensor>();
29+
copy_tensor(lod_tensor, &out_tensor);
30+
}
31+
32+
void operator()(const framework::LoDTensorArray &array) const {
33+
auto &out_array = *out_->GetMutable<framework::LoDTensorArray>();
34+
out_array.resize(array.size());
35+
for (size_t i = 0; i < array.size(); ++i) {
36+
copy_tensor(array[i], &out_array[i]);
37+
}
38+
}
39+
40+
void operator()(const framework::SelectedRows &rows) const {
41+
framework::SelectedRows &out_rows =
42+
*out_->GetMutable<framework::SelectedRows>();
43+
out_rows.set_rows(rows.rows());
44+
out_rows.set_height(rows.height());
45+
auto &t = rows.value();
46+
out_rows.mutable_value()->CopyFrom(t, t.place(), dev_ctx_);
47+
}
48+
49+
template <typename T>
50+
void operator()(const T &v) const {
51+
PADDLE_THROW("Not support type for assign op %s", typeid(T).name());
52+
}
53+
54+
private:
55+
void copy_tensor(const framework::LoDTensor &lod_tensor,
56+
framework::LoDTensor *out) const {
57+
auto &out_tensor = *out;
58+
out_tensor.CopyFrom(lod_tensor, lod_tensor.place(), dev_ctx_);
59+
out_tensor.set_lod(lod_tensor.lod());
60+
}
61+
62+
framework::Variable *out_;
63+
const platform::DeviceContext &dev_ctx_;
64+
};
65+
66+
class AssignOp : public framework::OperatorBase {
67+
public:
68+
AssignOp(const std::string &type, const framework::VariableNameMap &inputs,
69+
const framework::VariableNameMap &outputs,
70+
const framework::AttributeMap &attrs)
71+
: OperatorBase(type, inputs, outputs, attrs) {}
72+
void Run(const framework::Scope &scope,
73+
const platform::DeviceContext &dev_ctx) const override {
74+
auto *x = scope.FindVar(Input("X"));
75+
if (x == nullptr) {
76+
return;
77+
}
78+
auto *out = scope.FindVar(Output("Out"));
79+
PADDLE_ENFORCE(
80+
out != nullptr,
81+
"The Output(Out) should not be null if the Input(X) is set.");
82+
framework::VisitVarType(*x, AssignFunctor(out, dev_ctx));
83+
}
84+
};
85+
86+
class AssignOpProtoMaker : public framework::OpProtoAndCheckerMaker {
87+
public:
88+
AssignOpProtoMaker(framework::OpProto *proto,
89+
framework::OpAttrChecker *op_checker)
90+
: OpProtoAndCheckerMaker(proto, op_checker) {
91+
AddInput("X",
92+
"(LoDTensor, SelectedRows or LoDTensorArray) The input variable "
93+
"could be LoDTensor, SelectedRows or LoDTensorArray.")
94+
.AsDispensable();
95+
AddOutput("Out",
96+
"(LoDTensor, SelectedRows or LoDTensorArray) The type of output "
97+
"is the same as input X.");
98+
AddComment(R"DOC(Assign Operator
99+
100+
Out = X, when type in [LoDTensor/SelectedRows/LoDTensorArray]
101+
raise error if the type is not listed above.
102+
)DOC");
103+
}
104+
};
105+
106+
class AssignInferShape : public framework::InferShapeBase {
107+
public:
108+
void operator()(framework::InferShapeContext *context) const override {
109+
if (context->HasInput("X")) {
110+
auto type = context->GetInputsVarType("X")[0];
111+
if (type == framework::VarDesc_VarType_SELECTED_ROWS ||
112+
type == framework::VarDesc_VarType_LOD_TENSOR) {
113+
context->SetOutputDim("Out", context->GetInputDim("X"));
114+
}
115+
}
116+
}
117+
};
118+
119+
class AssignGradMaker : public framework::SingleGradOpDescMaker {
120+
public:
121+
using framework::SingleGradOpDescMaker::SingleGradOpDescMaker;
122+
123+
protected:
124+
std::unique_ptr<framework::OpDescBind> Apply() const override {
125+
auto *op = new framework::OpDescBind();
126+
op->SetType("assign");
127+
op->SetInput("X", OutputGrad("Out"));
128+
op->SetOutput("Out", InputGrad("X"));
129+
return std::unique_ptr<framework::OpDescBind>(op);
130+
}
131+
};
132+
133+
} // namespace operators
134+
} // namespace paddle
135+
136+
namespace ops = paddle::operators;
137+
REGISTER_OPERATOR(assign, ops::AssignOp, ops::AssignGradMaker,
138+
ops::AssignInferShape, ops::AssignOpProtoMaker);

‎python/paddle/v2/framework/tests/test_assign_op.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
import op_test
2+
import numpy
3+
import unittest
4+
5+
6+
class TestAssignOp(op_test.OpTest):
7+
def setUp(self):
8+
self.op_type = "assign"
9+
x = numpy.random.random(size=(100, 10))
10+
self.inputs = {'X': x}
11+
self.outputs = {'Out': x}
12+
13+
def test_forward(self):
14+
self.check_output()
15+
16+
def test_backward(self):
17+
self.check_grad(['X'], 'Out')
18+
19+
20+
if __name__ == '__main__':
21+
unittest.main()

0 commit comments

Comments
 (0)