m53297601
/
mindspore
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

!5508 Add AdamApplyOneWithDecayAssign fusion pass

Browse Source 
Merge pull request !5508 from YuJianfeng/adam_assign
pull/5508/MERGE
mindspore-ci-bot 5 years ago committed by Gitee
parent bfe3b9d171 4b77f6b53c
commit 3048240f16
5 changed files with 418 additions and 100 deletions
Show Stats Download Patch File Download Diff File

154
mindspore/ccsrc/backend/optimizer/ascend/ir_fusion/adam_apply_one_with_decay_rule.cc
Unescape View File

@ -24,7 +24,8 @@
namespace mindspore {
namespace opt {
std::vector<AnfNodePtr> AdamApplyOneWithDecayRule::GetFusionNodeInputs(const EquivPtr &equiv) const {
std::vector<AnfNodePtr> AdamApplyOneWithDecayRule::GetFusionNodeInputs(const EquivPtr &equiv,
const AnfNodePtr &final_node) const {
MS_EXCEPTION_IF_NULL(equiv);
auto input0 = utils::cast<AnfNodePtr>((*equiv)[input0_]);
auto input1 = utils::cast<AnfNodePtr>((*equiv)[input1_]);
@ -37,7 +38,12 @@ std::vector<AnfNodePtr> AdamApplyOneWithDecayRule::GetFusionNodeInputs(const Equ
auto mul3_x = utils::cast<AnfNodePtr>((*equiv)[mul3_x_]);
auto mul4_x = utils::cast<AnfNodePtr>((*equiv)[mul4_x_]);
auto add2_y = utils::cast<AnfNodePtr>((*equiv)[add2_y_]);
auto prim = std::make_shared<Primitive>(kAdamApplyOneWithDecayOpName);
PrimitivePtr prim = nullptr;
if (AnfAlgo::CheckPrimitiveType(final_node, prim::kPrimDepend)) {
prim = std::make_shared<Primitive>(kAdamApplyOneWithDecayAssignOpName);
} else {
prim = std::make_shared<Primitive>(kAdamApplyOneWithDecayOpName);
}
return {NewValueNode(prim), input0, input1, input2, input3, input4, mul0_x, mul1_x, mul2_x, mul3_x, mul4_x, add2_y};
}
@ -141,18 +147,152 @@ const BaseRef AdamApplyOneWithDecayRuleCond5::DefinePattern() const {
return sub0;
}
const BaseRef AdamApplyOneWithDecayAssignRuleCond1::DefinePattern() const {
auto sqrt = std::make_shared<Primitive>(kSqrtOpName);
auto real_div = std::make_shared<Primitive>(kRealDivOpName);
VectorRef mul0({prim::kPrimMul, mul0_x_, input2_});
VectorRef mul1({prim::kPrimMul, mul1_x_, input0_});
VectorRef square0({prim::kPrimSquare, input0_});
VectorRef add0({add0_var_, mul0, mul1});
VectorRef mul2({prim::kPrimMul, mul2_x_, input1_});
VectorRef mul3({prim::kPrimMul, mul3_x_, square0});
VectorRef add1({add1_var_, mul2, mul3});
VectorRef sqrt0({sqrt, add1});
VectorRef add2({prim::kPrimTensorAdd, add2_y_, sqrt0});
VectorRef mul4({prim::kPrimMul, mul4_x_, input3_});
VectorRef real_div0({real_div, add0, add2});
VectorRef add3({prim::kPrimTensorAdd, mul4, real_div0});
VectorRef mul5({prim::kPrimMul, input4_, add3});
VectorRef sub0({sub0_var_, input3_, mul5});
VectorRef assign0 = VectorRef({prim::kPrimAssign, input3_, sub0});
VectorRef depend0 = VectorRef({prim::kPrimDepend, sub0, assign0});
VectorRef assign1 = VectorRef({prim::kPrimAssign, input2_, add0});
VectorRef depend1 = VectorRef({prim::kPrimDepend, depend0, assign1});
VectorRef assign2 = VectorRef({prim::kPrimAssign, input1_, add1});
return VectorRef({prim::kPrimDepend, depend1, assign2});
}
const BaseRef AdamApplyOneWithDecayAssignRuleCond2::DefinePattern() const {
auto sqrt = std::make_shared<Primitive>(kSqrtOpName);
auto real_div = std::make_shared<Primitive>(kRealDivOpName);
VectorRef mul0({prim::kPrimMul, input2_, mul0_x_});
VectorRef mul1({prim::kPrimMul, input0_, mul1_x_});
VectorRef square0({prim::kPrimSquare, input0_});
VectorRef add0({add0_var_, mul0, mul1});
VectorRef mul2({prim::kPrimMul, input1_, mul2_x_});
VectorRef mul3({prim::kPrimMul, mul3_x_, square0});
VectorRef add1({add1_var_, mul2, mul3});
VectorRef sqrt0({sqrt, add1});
VectorRef add2({prim::kPrimTensorAdd, sqrt0, add2_y_});
VectorRef mul4({prim::kPrimMul, input3_, mul4_x_});
VectorRef real_div0({real_div, add0, add2});
VectorRef add3({prim::kPrimTensorAdd, mul4, real_div0});
VectorRef mul5({prim::kPrimMul, add3, input4_});
VectorRef sub0({sub0_var_, input3_, mul5});
VectorRef assign0 = VectorRef({prim::kPrimAssign, input3_, sub0});
VectorRef depend0 = VectorRef({prim::kPrimDepend, sub0, assign0});
VectorRef assign1 = VectorRef({prim::kPrimAssign, input2_, add0});
VectorRef depend1 = VectorRef({prim::kPrimDepend, depend0, assign1});
VectorRef assign2 = VectorRef({prim::kPrimAssign, input1_, add1});
return VectorRef({prim::kPrimDepend, depend1, assign2});
}
const BaseRef AdamApplyOneWithDecayAssignRuleCond3::DefinePattern() const {
auto sqrt = std::make_shared<Primitive>(kSqrtOpName);
auto real_div = std::make_shared<Primitive>(kRealDivOpName);
VectorRef mul0({prim::kPrimMul, mul0_x_, input2_});
VectorRef mul1({prim::kPrimMul, mul1_x_, input0_});
VectorRef square0({prim::kPrimSquare, input0_});
VectorRef add0({add0_var_, mul0, mul1});
VectorRef mul2({prim::kPrimMul, mul2_x_, input1_});
VectorRef mul3({prim::kPrimMul, square0, mul3_x_});
VectorRef add1({add1_var_, mul2, mul3});
VectorRef sqrt0({sqrt, add1});
VectorRef add2({prim::kPrimTensorAdd, sqrt0, add2_y_});
VectorRef mul4({prim::kPrimMul, mul4_x_, input3_});
VectorRef real_div0({real_div, add0, add2});
VectorRef add3({prim::kPrimTensorAdd, mul4, real_div0});
VectorRef mul5({prim::kPrimMul, add3, input4_});
VectorRef sub0({sub0_var_, input3_, mul5});
VectorRef assign0 = VectorRef({prim::kPrimAssign, input3_, sub0});
VectorRef depend0 = VectorRef({prim::kPrimDepend, sub0, assign0});
VectorRef assign1 = VectorRef({prim::kPrimAssign, input2_, add0});
VectorRef depend1 = VectorRef({prim::kPrimDepend, depend0, assign1});
VectorRef assign2 = VectorRef({prim::kPrimAssign, input1_, add1});
return VectorRef({prim::kPrimDepend, depend1, assign2});
}
const BaseRef AdamApplyOneWithDecayAssignRuleCond4::DefinePattern() const {
auto sqrt = std::make_shared<Primitive>(kSqrtOpName);
auto real_div = std::make_shared<Primitive>(kRealDivOpName);
VectorRef mul0({prim::kPrimMul, mul0_x_, input2_});
VectorRef mul1({prim::kPrimMul, mul1_x_, input0_});
VectorRef square0({prim::kPrimSquare, input0_});
VectorRef add0({add0_var_, mul0, mul1});
VectorRef mul2({prim::kPrimMul, mul2_x_, input1_});
VectorRef mul3({prim::kPrimMul, mul3_x_, square0});
VectorRef add1({add1_var_, mul2, mul3});
VectorRef sqrt0({sqrt, add1});
VectorRef add2({prim::kPrimTensorAdd, add2_y_, sqrt0});
VectorRef mul4({prim::kPrimMul, mul4_x_, input3_});
VectorRef real_div0({real_div, add0, add2});
VectorRef add3({prim::kPrimTensorAdd, mul4, real_div0});
VectorRef mul5({prim::kPrimMul, add3, input4_});
VectorRef sub0({sub0_var_, input3_, mul5});
VectorRef assign0 = VectorRef({prim::kPrimAssign, input3_, sub0});
VectorRef depend0 = VectorRef({prim::kPrimDepend, sub0, assign0});
VectorRef assign1 = VectorRef({prim::kPrimAssign, input2_, add0});
VectorRef depend1 = VectorRef({prim::kPrimDepend, depend0, assign1});
VectorRef assign2 = VectorRef({prim::kPrimAssign, input1_, add1});
return VectorRef({prim::kPrimDepend, depend1, assign2});
}
const BaseRef AdamApplyOneWithDecayAssignRuleCond5::DefinePattern() const {
auto sqrt = std::make_shared<Primitive>(kSqrtOpName);
auto real_div = std::make_shared<Primitive>(kRealDivOpName);
VectorRef mul0({prim::kPrimMul, mul0_x_, input2_});
VectorRef mul1({prim::kPrimMul, mul1_x_, input0_});
VectorRef square0({prim::kPrimSquare, input0_});
VectorRef add0({add0_var_, mul0, mul1});
VectorRef mul2({prim::kPrimMul, mul2_x_, input1_});
VectorRef mul3({prim::kPrimMul, mul3_x_, square0});
VectorRef add1({add1_var_, mul2, mul3});
VectorRef sqrt0({sqrt, add1});
VectorRef add2({prim::kPrimTensorAdd, sqrt0, add2_y_});
VectorRef mul4({prim::kPrimMul, mul4_x_, input3_});
VectorRef real_div0({real_div, add0, add2});
VectorRef add3({prim::kPrimTensorAdd, mul4, real_div0});
VectorRef mul5({prim::kPrimMul, add3, input4_});
VectorRef sub0({sub0_var_, input3_, mul5});
VectorRef assign0 = VectorRef({prim::kPrimAssign, input3_, sub0});
VectorRef depend0 = VectorRef({prim::kPrimDepend, sub0, assign0});
VectorRef assign1 = VectorRef({prim::kPrimAssign, input2_, add0});
VectorRef depend1 = VectorRef({prim::kPrimDepend, depend0, assign1});
VectorRef assign2 = VectorRef({prim::kPrimAssign, input1_, add1});
return VectorRef({prim::kPrimDepend, depend1, assign2});
}
const AnfNodePtr AdamApplyOneWithDecayRule::Process(const FuncGraphPtr &graph, const AnfNodePtr &node,
const EquivPtr &equiv) const {
if (graph == nullptr || node == nullptr || equiv == nullptr) {
return nullptr;
}
if (!CheckSupportDataType(node, kFloatDataTypeSet)) {
auto sub0 = node;
if (AnfAlgo::CheckPrimitiveType(node, prim::kPrimDepend)) {
auto iter_sub0 = (*equiv).find(sub0_var_);
if (iter_sub0 == (*equiv).end()) {
MS_LOG(EXCEPTION) << "The equiv map is expected to contains the sub0 var after matched.";
}
sub0 = utils::cast<AnfNodePtr>(iter_sub0->second);
}
MS_EXCEPTION_IF_NULL(sub0);
if (!CheckSupportDataType(sub0, kFloatDataTypeSet)) {
return nullptr;
}
std::vector<AnfNodePtr> inputs = GetFusionNodeInputs(equiv);
std::vector<AnfNodePtr> inputs = GetFusionNodeInputs(equiv, node);
auto fusion_node = graph->NewCNode(inputs);
MS_EXCEPTION_IF_NULL(fusion_node);
fusion_node->set_scope(node->scope());
fusion_node->set_scope(sub0->scope());
auto iter_add0 = (*equiv).find(add0_var_);
if (iter_add0 == (*equiv).end()) {
@ -167,9 +307,9 @@ const AnfNodePtr AdamApplyOneWithDecayRule::Process(const FuncGraphPtr &graph, c
auto add1 = utils::cast<AnfNodePtr>(iter_add1->second);
MS_EXCEPTION_IF_NULL(add1);
auto types = {AnfAlgo::GetOutputInferDataType(add1, 0), AnfAlgo::GetOutputInferDataType(add0, 0),
AnfAlgo::GetOutputInferDataType(node, 0)};
AnfAlgo::GetOutputInferDataType(sub0, 0)};
auto shapes = {AnfAlgo::GetOutputInferShape(add1, 0), AnfAlgo::GetOutputInferShape(add0, 0),
AnfAlgo::GetOutputInferShape(node, 0)};
AnfAlgo::GetOutputInferShape(sub0, 0)};
AnfAlgo::SetOutputInferTypeAndShape(types, shapes, fusion_node.get());
std::vector<AnfNodePtr> fusion_node_outputs;

49
mindspore/ccsrc/backend/optimizer/ascend/ir_fusion/adam_apply_one_with_decay_rule.h
Unescape View File

@ -40,13 +40,14 @@ class AdamApplyOneWithDecayRule : public PatternProcessPass {
add2_y_ = std::make_shared<Var>();
add0_var_ = std::make_shared<Var>(std::make_shared<Primitive>(prim::kPrimTensorAdd->name()));
add1_var_ = std::make_shared<Var>(std::make_shared<Primitive>(prim::kPrimTensorAdd->name()));
sub0_var_ = std::make_shared<Var>(std::make_shared<Primitive>(prim::kPrimSub->name()));
}
~AdamApplyOneWithDecayRule() override = default;
const BaseRef DefinePattern() const override = 0;
const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override;
protected:
std::vector<AnfNodePtr> GetFusionNodeInputs(const EquivPtr &equiv) const;
std::vector<AnfNodePtr> GetFusionNodeInputs(const EquivPtr &equiv, const AnfNodePtr &final_node) const;
VarPtr input0_;
VarPtr input1_;
VarPtr input2_;
@ -60,6 +61,7 @@ class AdamApplyOneWithDecayRule : public PatternProcessPass {
VarPtr add2_y_;
VarPtr add0_var_;
VarPtr add1_var_;
VarPtr sub0_var_;
};
class AdamApplyOneWithDecayRuleCond1 : public AdamApplyOneWithDecayRule {
@ -106,6 +108,51 @@ class AdamApplyOneWithDecayRuleCond5 : public AdamApplyOneWithDecayRule {
~AdamApplyOneWithDecayRuleCond5() override = default;
const BaseRef DefinePattern() const override;
};
class AdamApplyOneWithDecayAssignRuleCond1 : public AdamApplyOneWithDecayRule {
public:
explicit AdamApplyOneWithDecayAssignRuleCond1(bool multigraph = true)
: AdamApplyOneWithDecayRule("adam_apply_one_with_decay_assign_rule_cond1", multigraph) {}
~AdamApplyOneWithDecayAssignRuleCond1() override = default;
const BaseRef DefinePattern() const override;
};
class AdamApplyOneWithDecayAssignRuleCond2 : public AdamApplyOneWithDecayRule {
public:
explicit AdamApplyOneWithDecayAssignRuleCond2(bool multigraph = true)
: AdamApplyOneWithDecayRule("adam_apply_one_with_decay_assign_rule_cond2", multigraph) {}
~AdamApplyOneWithDecayAssignRuleCond2() override = default;
const BaseRef DefinePattern() const override;
};
class AdamApplyOneWithDecayAssignRuleCond3 : public AdamApplyOneWithDecayRule {
public:
explicit AdamApplyOneWithDecayAssignRuleCond3(bool multigraph = true)
: AdamApplyOneWithDecayRule("adam_apply_one_with_decay_assign_rule_cond3", multigraph) {}
~AdamApplyOneWithDecayAssignRuleCond3() override = default;
const BaseRef DefinePattern() const override;
};
class AdamApplyOneWithDecayAssignRuleCond4 : public AdamApplyOneWithDecayRule {
public:
explicit AdamApplyOneWithDecayAssignRuleCond4(bool multigraph = true)
: AdamApplyOneWithDecayRule("adam_apply_one_with_decay_assign_rule_cond4", multigraph) {}
~AdamApplyOneWithDecayAssignRuleCond4() override = default;
const BaseRef DefinePattern() const override;
};
class AdamApplyOneWithDecayAssignRuleCond5 : public AdamApplyOneWithDecayRule {
public:
explicit AdamApplyOneWithDecayAssignRuleCond5(bool multigraph = true)
: AdamApplyOneWithDecayRule("adam_apply_one_with_decay_assign_rule_cond5", multigraph) {}
~AdamApplyOneWithDecayAssignRuleCond5() override = default;
const BaseRef DefinePattern() const override;
};
} // namespace opt
} // namespace mindspore
#endif // MINDSPORE_CCSRC_BACKEND_OPTIMIZER_ASCEND_IR_FUSION_ADAM_APPLY_ONE_WITH_DECAY_RULE_H_

1
mindspore/ccsrc/utils/utils.h
Unescape View File

@ -122,6 +122,7 @@ constexpr auto kLayerNormBetaGammaBackpropOpName = "LayerNormBetaGammaBackprop";
constexpr auto kLambNextMVOpName = "LambNextMV";
constexpr auto kConfusionTransposeDOpName = "ConfusionTransposeD";
constexpr auto kAdamApplyOneWithDecayOpName = "AdamApplyOneWithDecay";
constexpr auto kAdamApplyOneWithDecayAssignOpName = "AdamApplyOneWithDecayAssign";
constexpr auto kBatchNormGradOpName = "BatchNormGrad";
constexpr auto kBNInferOpName = "BNInfer";
constexpr auto kAdamApplyOneOpName = "AdamApplyOne";

125
tests/ut/cpp/pre_activate/ascend/ir_fusion/adam_apply_one_with_decay_rule_test.cc
Unescape View File

@ -31,7 +31,7 @@ class TestHWOptimizeAdamApplyOneWithDecayRule : public BackendCommon {
};
TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_rule_cond1) {
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule_cond1", "before");
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule", "before_cond1");
std::vector<int> shp{2, 32, 224, 224};
auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp);
@ -47,12 +47,12 @@ TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_r
optimizer->AddPassManager(pm);
FuncGraphPtr new_graph = optimizer->Optimize(fg);
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule_cond1", "after");
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule", "after");
EXPECT_TRUE(CheckEqualGraph(g_after, new_graph));
}
TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_rule_cond2) {
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule_cond2", "before");
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule", "before_cond2");
std::vector<int> shp{2, 32, 224, 224};
auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp);
@ -68,12 +68,12 @@ TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_r
optimizer->AddPassManager(pm);
FuncGraphPtr new_graph = optimizer->Optimize(fg);
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule_cond2", "after");
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule", "after");
EXPECT_TRUE(CheckEqualGraph(g_after, new_graph));
}
TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_rule_cond3) {
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule_cond3", "before");
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule", "before_cond3");
std::vector<int> shp{2, 32, 224, 224};
auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp);
@ -89,12 +89,12 @@ TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_r
optimizer->AddPassManager(pm);
FuncGraphPtr new_graph = optimizer->Optimize(fg);
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule_cond3", "after");
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule", "after");
EXPECT_TRUE(CheckEqualGraph(g_after, new_graph));
}
TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_rule_cond4) {
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule_cond4", "before");
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule", "before_cond4");
std::vector<int> shp{2, 32, 224, 224};
auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp);
@ -110,12 +110,12 @@ TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_r
optimizer->AddPassManager(pm);
FuncGraphPtr new_graph = optimizer->Optimize(fg);
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule_cond4", "after");
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule", "after");
EXPECT_TRUE(CheckEqualGraph(g_after, new_graph));
}
TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_rule_cond5) {
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule_cond5", "before");
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule", "before_cond5");
std::vector<int> shp{2, 32, 224, 224};
auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp);
@ -131,7 +131,112 @@ TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_r
optimizer->AddPassManager(pm);
FuncGraphPtr new_graph = optimizer->Optimize(fg);
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule_cond5", "after");
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_rule", "after");
EXPECT_TRUE(CheckEqualGraph(g_after, new_graph));
}
TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_assign_rule_cond1) {
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_assign_rule", "before_cond1");
std::vector<int> shp{2, 32, 224, 224};
auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp);
AbstractBasePtrList args_spec_list;
for (size_t i = 0; i < 11; ++i) {
args_spec_list.push_back(x_abstract);
}
auto fg = GetKernelGraph(g, args_spec_list);
auto optimizer = std::make_shared<opt::GraphOptimizer>();
auto pm = std::make_shared<opt::PassManager>();
pm->AddPass(std::make_shared<opt::AdamApplyOneWithDecayAssignRuleCond1>());
optimizer->AddPassManager(pm);
FuncGraphPtr new_graph = optimizer->Optimize(fg);
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_assign_rule", "after");
EXPECT_TRUE(CheckEqualGraph(g_after, new_graph));
}
TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_assign_rule_cond2) {
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_assign_rule", "before_cond2");
std::vector<int> shp{2, 32, 224, 224};
auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp);
AbstractBasePtrList args_spec_list;
for (size_t i = 0; i < 11; ++i) {
args_spec_list.push_back(x_abstract);
}
auto fg = GetKernelGraph(g, args_spec_list);
auto optimizer = std::make_shared<opt::GraphOptimizer>();
auto pm = std::make_shared<opt::PassManager>();
pm->AddPass(std::make_shared<opt::AdamApplyOneWithDecayAssignRuleCond2>());
optimizer->AddPassManager(pm);
FuncGraphPtr new_graph = optimizer->Optimize(fg);
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_assign_rule", "after");
EXPECT_TRUE(CheckEqualGraph(g_after, new_graph));
}
TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_assign_rule_cond3) {
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_assign_rule", "before_cond3");
std::vector<int> shp{2, 32, 224, 224};
auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp);
AbstractBasePtrList args_spec_list;
for (size_t i = 0; i < 11; ++i) {
args_spec_list.push_back(x_abstract);
}
auto fg = GetKernelGraph(g, args_spec_list);
auto optimizer = std::make_shared<opt::GraphOptimizer>();
auto pm = std::make_shared<opt::PassManager>();
pm->AddPass(std::make_shared<opt::AdamApplyOneWithDecayAssignRuleCond3>());
optimizer->AddPassManager(pm);
FuncGraphPtr new_graph = optimizer->Optimize(fg);
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_assign_rule", "after");
EXPECT_TRUE(CheckEqualGraph(g_after, new_graph));
}
TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_assign_rule_cond4) {
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_assign_rule", "before_cond4");
std::vector<int> shp{2, 32, 224, 224};
auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp);
AbstractBasePtrList args_spec_list;
for (size_t i = 0; i < 11; ++i) {
args_spec_list.push_back(x_abstract);
}
auto fg = GetKernelGraph(g, args_spec_list);
auto optimizer = std::make_shared<opt::GraphOptimizer>();
auto pm = std::make_shared<opt::PassManager>();
pm->AddPass(std::make_shared<opt::AdamApplyOneWithDecayAssignRuleCond4>());
optimizer->AddPassManager(pm);
FuncGraphPtr new_graph = optimizer->Optimize(fg);
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_assign_rule", "after");
EXPECT_TRUE(CheckEqualGraph(g_after, new_graph));
}
TEST_F(TestHWOptimizeAdamApplyOneWithDecayRule, test_adam_apply_one_with_decay_assign_rule_cond5) {
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_assign_rule", "before_cond5");
std::vector<int> shp{2, 32, 224, 224};
auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp);
AbstractBasePtrList args_spec_list;
for (size_t i = 0; i < 11; ++i) {
args_spec_list.push_back(x_abstract);
}
auto fg = GetKernelGraph(g, args_spec_list);
auto optimizer = std::make_shared<opt::GraphOptimizer>();
auto pm = std::make_shared<opt::PassManager>();
pm->AddPass(std::make_shared<opt::AdamApplyOneWithDecayAssignRuleCond5>());
optimizer->AddPassManager(pm);
FuncGraphPtr new_graph = optimizer->Optimize(fg);
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_adam_apply_one_with_decay_assign_rule", "after");
EXPECT_TRUE(CheckEqualGraph(g_after, new_graph));
}
} // namespace opt

189
tests/ut/cpp/python_input/gtest_input/pre_activate/adam_apply_one_with_decay_rule.py
View File

File diff suppressed because it is too large Load Diff
Write Preview
Loading…
Cancel
Save