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

Add BatchNormGrad split pass

Browse Source
pull/823/head
huanghui 5 years ago
parent 6241814320
commit 0de0aee54e
5 changed files with 287 additions and 0 deletions
Show Stats Download Patch File Download Diff File

132
mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_grad_split.cc
Unescape View File

@ -0,0 +1,132 @@
/**
* Copyright 2020 Huawei Technologies Co., Ltd
*
* 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 "pre_activate/ascend/ir_fission/batch_norm_grad_split.h"
#include <vector>
#include <string>
#include <memory>
#include "utils/utils.h"
#include "utils/context/ms_context.h"
#include "common/utils.h"
#include "pre_activate/common/helper.h"
#include "device/kernel_info.h"
#include "session/anf_runtime_algorithm.h"
namespace mindspore {
namespace opt {
namespace {
void CreateOutputsOfUpdateGrad(const FuncGraphPtr &graph, const CNodePtr &bn_grad_node,
std::vector<AnfNodePtr> *bn_update_grad_outputs) {
MS_EXCEPTION_IF_NULL(graph);
MS_EXCEPTION_IF_NULL(bn_grad_node);
auto bn_grad_inputs = bn_grad_node->inputs();
if (bn_grad_inputs.size() < kBNGradInputNum) {
MS_LOG(EXCEPTION) << "BNGrad has wrong inputs size";
}
std::vector<AnfNodePtr> bn_update_grad_inputs = {
NewValueNode(std::make_shared<Primitive>(kBNTrainingUpdateGradOpName)), bn_grad_inputs[1], bn_grad_inputs[2],
bn_grad_inputs[4], bn_grad_inputs[5]};
auto bn_update_grad = graph->NewCNode(bn_update_grad_inputs);
MS_EXCEPTION_IF_NULL(bn_update_grad);
bn_update_grad->set_kernel_info(std::make_shared<device::KernelInfo>());
bn_update_grad->set_scope(bn_grad_node->scope());
auto types = {AnfAlgo::GetOutputInferDataType(bn_grad_node, 1), AnfAlgo::GetOutputInferDataType(bn_grad_node, 2)};
auto shapes = {AnfAlgo::GetOutputInferShape(bn_grad_node, 1), AnfAlgo::GetOutputInferShape(bn_grad_node, 2)};
AnfAlgo::SetOutputInferTypeAndShape(types, shapes, bn_update_grad.get());
AnfAlgo::CopyNodeAttr(kAttrEpsilon, bn_grad_node, bn_update_grad);
CreateMultipleOutputsOfAnfNode(graph, bn_update_grad, kBNTrainingUpdateGradOutputNum, bn_update_grad_outputs);
}
void CreateOutputsOfReduceGrad(const FuncGraphPtr &graph, const CNodePtr &bn_grad_node,
const std::vector<AnfNodePtr> &bn_update_grad_outputs,
std::vector<AnfNodePtr> *bn_reduce_grad_outputs) {
MS_EXCEPTION_IF_NULL(graph);
MS_EXCEPTION_IF_NULL(bn_grad_node);
auto bn_grad_inputs = bn_grad_node->inputs();
if (bn_grad_inputs.size() < kBNGradInputNum) {
MS_LOG(EXCEPTION) << "BNGrad has wrong inputs size";
}
if (bn_update_grad_outputs.size() != kBNTrainingUpdateGradOutputNum) {
MS_LOG(EXCEPTION) << "BNTrainingReduceGrad_outputs has wrong size";
}
std::vector<AnfNodePtr> bn_reduce_grad_inputs = {
NewValueNode(std::make_shared<Primitive>(kBNTrainingReduceGradOpName)),
bn_grad_inputs[1],
bn_grad_inputs[2],
bn_update_grad_outputs[0],
bn_update_grad_outputs[1],
bn_grad_inputs[3],
bn_grad_inputs[4],
bn_grad_inputs[5]};
auto bn_reduce_grad = graph->NewCNode(bn_reduce_grad_inputs);
MS_EXCEPTION_IF_NULL(bn_reduce_grad);
bn_reduce_grad->set_kernel_info(std::make_shared<device::KernelInfo>());
bn_reduce_grad->set_scope(bn_grad_node->scope());
auto types = {AnfAlgo::GetOutputInferDataType(bn_grad_node, 0)};
auto shapes = {AnfAlgo::GetOutputInferShape(bn_grad_node, 0)};
AnfAlgo::SetOutputInferTypeAndShape(types, shapes, bn_reduce_grad.get());
AnfAlgo::CopyNodeAttr(kAttrEpsilon, bn_grad_node, bn_reduce_grad);
(*bn_reduce_grad_outputs).push_back(bn_reduce_grad);
}
} // namespace
const BaseRef BatchNormGradSplit::DefinePattern() const {
VarPtr Xs = std::make_shared<SeqVar>();
auto prim = std::make_shared<Primitive>(kBatchNormGradOpName);
return VectorRef({prim, Xs});
}
const AnfNodePtr BatchNormGradSplit::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
const EquivPtr &) const {
MS_EXCEPTION_IF_NULL(node);
MS_EXCEPTION_IF_NULL(func_graph);
auto cnode = node->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(cnode);
auto primitive = AnfAlgo::GetCNodePrimitive(cnode);
MS_EXCEPTION_IF_NULL(primitive);
if (!primitive->HasAttr(kAttrIsTraining)) {
MS_LOG(INFO) << "Op BatchNormGrad must have attrs of is_training";
return nullptr;
}
if (!AnfAlgo::GetNodeAttr<bool>(cnode, kAttrIsTraining)) {
MS_LOG(INFO) << "is_training must be true";
return nullptr;
}
std::vector<AnfNodePtr> bn_update_grad_outputs;
CreateOutputsOfUpdateGrad(func_graph, cnode, &bn_update_grad_outputs);
if (bn_update_grad_outputs.size() != kBNTrainingUpdateGradOutputNum) {
MS_LOG(EXCEPTION) << "bn_update_grad_outputs has wrong size";
}
std::vector<AnfNodePtr> bn_reduce_grad_outputs;
CreateOutputsOfReduceGrad(func_graph, cnode, bn_update_grad_outputs, &bn_reduce_grad_outputs);
if (bn_reduce_grad_outputs.size() != kSingleOutputNum) {
MS_LOG(EXCEPTION) << "bn_reduce_grad_outputs has wrong size";
}
std::vector<AnfNodePtr> make_tuple_inputs = {NewValueNode(prim::kPrimMakeTuple), bn_reduce_grad_outputs[0],
bn_update_grad_outputs[0], bn_update_grad_outputs[1]};
auto make_tuple = func_graph->NewCNode(make_tuple_inputs);
return make_tuple;
}
} // namespace opt
} // namespace mindspore

33
mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_grad_split.h
Unescape View File

@ -0,0 +1,33 @@
/**
* Copyright 2020 Huawei Technologies Co., Ltd
*
* 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.
*/
#ifndef MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_IR_FISSION_BATCH_NORM_GRAD_SPLIT_H_
#define MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_IR_FISSION_BATCH_NORM_GRAD_SPLIT_H_
#include "pre_activate/common/optimizer.h"
#include "pre_activate/common/helper.h"
namespace mindspore {
namespace opt {
class BatchNormGradSplit : public PatternProcessPass {
public:
explicit BatchNormGradSplit(bool multigraph = true) : PatternProcessPass("batch_norm_grad_split", multigraph) {}
~BatchNormGradSplit() override = default;
const BaseRef DefinePattern() const override;
const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override;
};
} // namespace opt
} // namespace mindspore
#endif // MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_IR_FISSION_BATCH_NORM_GRAD_SPLIT_H_

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

@ -107,6 +107,7 @@ constexpr auto kLambNextMVOpName = "LambNextMV";
constexpr auto kConfusionTransposeDOpName = "ConfusionTransposeD";
constexpr auto kAdamApplyOneWithDecayOpName = "AdamApplyOneWithDecay";
constexpr auto kBatchNormOpName = "BatchNorm";
constexpr auto kBatchNormGradOpName = "BatchNormGrad";
constexpr auto kAdamApplyOneOpName = "AdamApplyOne";
constexpr auto kDropoutGenMask = "DropoutGenMask";
constexpr auto kResizeNearestNeighborGrad = "ResizeNearestNeighborGrad";
@ -162,6 +163,7 @@ constexpr auto kAttrLabelForInsertStreamActive = "label_for_insert_stream_active
constexpr auto kAttrFusion = "fusion";
constexpr auto kAttrGroup = "group";
constexpr auto kAttrOp = "op";
constexpr auto kAttrIsTraining = "is_training";
// attr value
constexpr auto kValueTargetSwitch = "target_switch";

59
tests/ut/cpp/pre_activate/ascend/ir_fission/batch_norm_grad_split_test.cc
Unescape View File

@ -0,0 +1,59 @@
/**
* Copyright 2020 Huawei Technologies Co., Ltd
*
* 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 "common/backend_common_test.h"
#include "common/py_func_graph_fetcher.h"
#include "operator/ops.h"
#include "ir/meta_tensor.h"
#include "debug/anf_ir_dump.h"
#include "utils/utils.h"
#include "pre_activate/common/optimizer.h"
#include "pre_activate/ascend/ir_fission/batch_norm_grad_split.h"
#include "session/anf_runtime_algorithm.h"
namespace mindspore {
namespace opt {
class TestHWBatchNormGradSplit : public BackendCommon {
public:
TestHWBatchNormGradSplit() : get_py_fun_("gtest_input.pre_activate.batch_norm_grad_split", true) {}
public:
UT::PyFuncGraphFetcher get_py_fun_;
};
TEST_F(TestHWBatchNormGradSplit, test_split) {
get_py_fun_.SetDoResolve(true);
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_batch_norm_grad_split", "before");
EXPECT_NE(g, nullptr);
std::vector<int> shp_x{1, 64, 112, 112};
std::vector<int> shp_b{64};
auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp_x);
auto b_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp_b);
AbstractBasePtrList args_spec_list{x_abstract, x_abstract, b_abstract, b_abstract, b_abstract, b_abstract};
auto kernel_graph = GetKernelGraph(g, args_spec_list);
EXPECT_NE(kernel_graph, nullptr);
auto optimizer = std::make_shared<opt::GraphOptimizer>();
auto pm = std::make_shared<opt::PassManager>();
auto pass = std::make_shared<opt::BatchNormGradSplit>();
pm->AddPass(pass);
optimizer->AddPassManager(pm);
auto new_graph = optimizer->Optimize(kernel_graph);
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_batch_norm_grad_split", "after");
EXPECT_TRUE(CheckEqualGraph(g_after, new_graph));
}
} // namespace opt
} // namespace mindspore

61
tests/ut/cpp/python_input/gtest_input/pre_activate/batch_norm_grad_split.py
Unescape View File

@ -0,0 +1,61 @@
# Copyright 2020 Huawei Technologies Co., Ltd
#
# 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.
# ============================================================================
from mindspore.ops import operations as P
from mindspore.ops.operations import _grad_ops as G
from mindspore.ops import Primitive
batch_norm_grad = G.BatchNormGrad(is_training=True)
bn_training_update_grad = Primitive('BNTrainingUpdateGrad')
bn_training_reduce_grad = Primitive('BNTrainingReduceGrad')
make_tuple = Primitive('make_tuple')
tuple_getitem = Primitive('tuple_getitem')
class FnDict:
def __init__(self):
self.fnDict = {}
def __call__(self, fn):
self.fnDict[fn.__name__] = fn
def __getitem__(self, name):
return self.fnDict[name]
def test_batch_norm_grad_split(tag):
fns = FnDict()
@fns
def before(i0, i1, i2, i3, i4, i5):
bn_grad_output = batch_norm_grad(i0, i1, i2, i3, i4, i5)
item0 = tuple_getitem(bn_grad_output, 0)
item1 = tuple_getitem(bn_grad_output, 1)
item2 = tuple_getitem(bn_grad_output, 2)
output = make_tuple(item0, item1, item2)
return output
@fns
def after(i0, i1, i2, i3, i4, i5):
bn_update_grad_output = bn_training_update_grad(i0, i1, i3, i4)
update_item0 = tuple_getitem(bn_update_grad_output, 0)
update_item1 = tuple_getitem(bn_update_grad_output, 1)
bn_reduce_grad_output = bn_training_reduce_grad(i0, i1, update_item0, update_item1, i2, i3, i4)
output = make_tuple(bn_reduce_grad_output, update_item0, update_item1)
item0 = tuple_getitem(output, 0)
item1 = tuple_getitem(output, 1)
item2 = tuple_getitem(output, 2)
output = make_tuple(item0, item1, item2)
return make_tuple(output)
return fns[tag]
Write Preview
Loading…
Cancel
Save