From ae9d447fca8869f63dca03cb778805299ce48afb Mon Sep 17 00:00:00 2001 From: qiaolongfei Date: Thu, 13 Jul 2017 11:36:55 +0800 Subject: [PATCH] Add Init to OperatorBase --- paddle/framework/op_registry.h | 1 + paddle/framework/operator.h | 4 ++++ paddle/framework/operator_test.cc | 5 +++++ 3 files changed, 10 insertions(+) diff --git a/paddle/framework/op_registry.h b/paddle/framework/op_registry.h index 248c7a1a3b866..e46da822c6c2e 100644 --- a/paddle/framework/op_registry.h +++ b/paddle/framework/op_registry.h @@ -119,6 +119,7 @@ class OpRegistry { op->attrs_[attr.name()] = AttrTypeHelper::GetAttrValue(attr); } op_checkers().at(op_type).Check(op->attrs_); + op->Init(); return op; } diff --git a/paddle/framework/operator.h b/paddle/framework/operator.h index 0ce422e007c39..4336115670031 100644 --- a/paddle/framework/operator.h +++ b/paddle/framework/operator.h @@ -49,6 +49,10 @@ class OperatorBase { std::string DebugString() const; + /// Init will be called after CreateOperator, you can put some initialization + /// logic here. + virtual void Init() {} + /// InferShape infer the size of Variables used by this Operator with /// information inside scope virtual void InferShape(const std::shared_ptr& scope) const = 0; diff --git a/paddle/framework/operator_test.cc b/paddle/framework/operator_test.cc index be8c4be2d4296..01b87bb50e7a1 100644 --- a/paddle/framework/operator_test.cc +++ b/paddle/framework/operator_test.cc @@ -21,14 +21,19 @@ namespace framework { class OperatorTest : public OperatorBase { public: + void Init() override { x = 1; } void InferShape(const std::shared_ptr& scope) const override {} void Run(const std::shared_ptr& scope, const platform::DeviceContext& dev_ctx) const override { float scale = GetAttr("scale"); ASSERT_NEAR(scale, 3.14, 1e-5); ASSERT_EQ(scope->GetVariable(inputs_[0]), nullptr); + ASSERT_EQ(x, 1); ASSERT_NE(scope->GetVariable(outputs_[0]), nullptr); } + + public: + float x = 0; }; class OperatorTestProtoAndCheckerMaker : public OpProtoAndCheckerMaker {