Huawei_Technology
/
mindspore
Not watched
1
0
Fork 0
Code
Releases 13 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
Browse Source

!11537 layer norm fusion support new IR 

From: @wangzhe128
Reviewed-by: 
Signed-off-by:
tags/v1.2.0-rc1
mindspore-ci-bot Gitee 5 years ago
parent
48f2d82c0e e4e5c1231c
commit
e2fc667481
7 changed files with 128 additions and 185 deletions
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. +1
    -1
      mindspore/lite/src/ops/layer_norm.cc
  2. +1
    -0
      mindspore/lite/src/runtime/agent/npu/CMakeLists.txt
  3. +0
    -3
      mindspore/lite/src/runtime/kernel/arm/fp32/layer_norm_fp32.cc
  4. +0
    -3
      mindspore/lite/src/runtime/kernel/arm/int8/layer_norm_int8.cc
  5. +3
    -3
      mindspore/lite/tools/converter/anf_transform.cc
  6. +110
    -169
      mindspore/lite/tools/optimizer/fusion/layer_norm_fusion.cc
  7. +13
    -6
      mindspore/lite/tools/optimizer/fusion/layer_norm_fusion.h

+ 1
- 1
mindspore/lite/src/ops/layer_norm.cc View File

@@ -81,7 +81,7 @@ int LayerNorm::UnPackToFlatBuilder(const schema::Primitive *primitive, flatbuffe
return RET_ERROR;
}

auto val_offset = schema::CreateLayerNorm(*fbb, attr->epsilon(), attr->begin_norm_axis(), attr->begin_params_axis());
auto val_offset = schema::CreateLayerNorm(*fbb, attr->begin_norm_axis(), attr->begin_params_axis(), attr->epsilon());
auto prim_offset = schema::CreatePrimitive(*fbb, schema::PrimitiveType_LayerNorm, val_offset.o);
fbb->Finish(prim_offset);
return RET_OK;


+ 1
- 0
mindspore/lite/src/runtime/agent/npu/CMakeLists.txt View File

@@ -14,6 +14,7 @@ add_library(hiai_ir_build SHARED IMPORTED)
set_target_properties(hiai_ir_build PROPERTIES IMPORTED_LOCATION
${DDK_LIB_PATH}/libhiai_ir_build.so)
add_library(npu_kernel_mid OBJECT ${NPU_RUNTIME_SRC})
add_dependencies(npu_kernel_mid fbs_src)
target_link_libraries(
npu_kernel_mid
hiai


+ 0
- 3
mindspore/lite/src/runtime/kernel/arm/fp32/layer_norm_fp32.cc View File

@@ -34,9 +34,6 @@ int LayerNormCPUKernel::Init() {
}

int LayerNormCPUKernel::ReSize() {
param_->begin_norm_axis_ = -1;
param_->begin_params_axis_ = -1;

auto shape = in_tensors_.front()->shape();
param_->begin_norm_axis_ =
param_->begin_norm_axis_ > 0 ? param_->begin_norm_axis_ : param_->begin_norm_axis_ + shape.size();


+ 0
- 3
mindspore/lite/src/runtime/kernel/arm/int8/layer_norm_int8.cc View File

@@ -82,9 +82,6 @@ int LayerNormInt8CPUKernel::Init() {
}

int LayerNormInt8CPUKernel::ReSize() {
param_->begin_norm_axis_ = -1;
param_->begin_params_axis_ = -1;

auto shape = in_tensors_.front()->shape();
param_->begin_norm_axis_ =
param_->begin_norm_axis_ > 0 ? param_->begin_norm_axis_ : param_->begin_norm_axis_ + shape.size();


+ 3
- 3
mindspore/lite/tools/converter/anf_transform.cc View File

@@ -121,9 +121,6 @@ int AnfTransform::AddGraphPass(const std::shared_ptr<opt::GraphOptimizer> &optim
weight_format_transform_pass->SetFmkType(config->fmk);
weight_format_transform_pass->SetQuantType(config->quantType);
graph_pm->AddPass(weight_format_transform_pass);
auto infershape_pass = std::make_shared<opt::InferShapePass>();
infershape_pass->SetFmkType(config->fmk);
graph_pm->AddPass(infershape_pass);
auto slice_prepose_pass = std::make_shared<opt::SlicePreposePass>();
slice_prepose_pass->SetFmkType(config->fmk);
graph_pm->AddPass(slice_prepose_pass);
@@ -155,6 +152,9 @@ int AnfTransform::AddConstFoldPass(const std::shared_ptr<opt::GraphOptimizer> &o
auto update_conv2d_param_pass = std::make_shared<opt::UpdateConv2DParamPass>();
update_conv2d_param_pass->SetFmkType(config->fmk);
const_fold_pm->AddPass(update_conv2d_param_pass);
auto infershape_pass = std::make_shared<opt::InferShapePass>();
infershape_pass->SetFmkType(config->fmk);
const_fold_pm->AddPass(infershape_pass);
optimizer->AddPassManager(const_fold_pm);
return RET_OK;
}


+ 110
- 169
mindspore/lite/tools/optimizer/fusion/layer_norm_fusion.cc View File

@@ -30,11 +30,6 @@
namespace mindspore {
namespace opt {
namespace {
constexpr size_t kAddInputsLength = 3;
constexpr size_t kSubInputsLength = 3;
constexpr size_t kMulInputsLength = 3;
constexpr size_t kRsqrtInputsLength = 2;
constexpr size_t kReduceInputsLength = 2;

bool IsAddNode(const BaseRef &n) {
if (utils::isa<CNodePtr>(n) || utils::isa<ValueNodePtr>(n)) {
@@ -52,14 +47,6 @@ bool IsSquaredDifferenceNode(const BaseRef &n) {
return false;
}

bool IsReduceNode(const BaseRef &n) {
if (utils::isa<CNodePtr>(n) || utils::isa<ValueNodePtr>(n)) {
auto type = opt::GetCNodeType(n);
return type == schema::PrimitiveType_Reduce;
}
return false;
}

bool IsRsqrtNode(const BaseRef &n) {
if (utils::isa<CNodePtr>(n) || utils::isa<ValueNodePtr>(n)) {
auto type = opt::GetCNodeType(n);
@@ -86,13 +73,11 @@ bool IsSubNode(const BaseRef &n) {
} // namespace

const BaseRef LayerNormFusion::DefinePattern() const {
auto mean1 = std::make_shared<CondVar>(IsReduceNode);
VectorRef mean1_ref = VectorRef({mean1, input_});
VectorRef mean1_ref = VectorRef({mean1_, input_});
auto squared_diffference1 = std::make_shared<CondVar>(IsSquaredDifferenceNode);
VectorRef squared_diffference1_ref = VectorRef({squared_diffference1, input_, mean1_ref});
auto mul1 = std::make_shared<CondVar>(IsMulNode);
auto mean2 = std::make_shared<CondVar>(IsReduceNode);
VectorRef mean2_ref = VectorRef({mean2, squared_diffference1_ref});
VectorRef mean2_ref = VectorRef({mean2_, squared_diffference1_ref});
auto add1 = std::make_shared<CondVar>(IsAddNode);
VectorRef add1_ref = VectorRef({add1, mean2_ref, epsilon_});
auto rsqrt1 = std::make_shared<CondVar>(IsRsqrtNode);
@@ -109,221 +94,177 @@ const BaseRef LayerNormFusion::DefinePattern() const {
return add2_ref;
}

CNodePtr LayerNormFusion::CreateLayerNormNode(const FuncGraphPtr &func_graph, const EquivPtr &equiv,
const std::vector<int> &shape, const float epsilon) const {
MS_EXCEPTION_IF_NULL(func_graph);
CNodePtr LayerNormFusion::CreateLayerNormNode(const FuncGraphPtr &func_graph, const EquivPtr &equiv, float epsilon,
int begin_norm_axis, int begin_params_axis) const {
MS_ASSERT(func_graph != nullptr);
MS_ASSERT(equiv != nullptr);
auto layer_norm_primitive = std::make_unique<schema::PrimitiveT>();
std::unique_ptr<schema::LayerNormT> attr = std::make_unique<schema::LayerNormT>();
attr->epsilon = epsilon;
attr->begin_norm_axis = begin_norm_axis;
attr->begin_params_axis = begin_params_axis;
layer_norm_primitive->value.type = schema::PrimitiveType_LayerNorm;
layer_norm_primitive->value.value = attr.release();
auto layer_norm_cvalue = lite::PrimitiveC::Create(layer_norm_primitive.release());
auto value_node = NewValueNode(std::shared_ptr<lite::PrimitiveC>(layer_norm_cvalue));
std::vector<AnfNodePtr> new_node_inputs = {value_node};
auto input_node = utils::cast<AnfNodePtr>((*equiv)[input_]);
MS_EXCEPTION_IF_NULL(input_node);
MS_ASSERT(input_node != nullptr);
new_node_inputs.push_back(input_node);
auto gamma_node = utils::cast<AnfNodePtr>((*equiv)[gamma_]);
MS_EXCEPTION_IF_NULL(gamma_node);
MS_ASSERT(gamma_node != nullptr);
new_node_inputs.push_back(gamma_node);
auto beta_node = utils::cast<AnfNodePtr>((*equiv)[beta_]);
MS_EXCEPTION_IF_NULL(beta_node);
MS_ASSERT(beta_node != nullptr);
new_node_inputs.push_back(beta_node);
auto new_node = func_graph->NewCNode(new_node_inputs);
return new_node;
}

const AnfNodePtr LayerNormFusion::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
const EquivPtr &equiv) const {
MS_ASSERT(func_graph != nullptr);
MS_ASSERT(node != nullptr);
MS_LOG(DEBUG) << "layer_norm pass";
if (CheckIfFuncGraphIsNull(func_graph) != lite::RET_OK || CheckIfAnfNodeIsNull(node) != lite::RET_OK) {
lite::ReturnCode::GetSingleReturnCode()->UpdateReturnCode(lite::RET_NULL_PTR);
return nullptr;
}

// add2
auto add2_cnode = node->cast<CNodePtr>();
if (CheckIfCNodeIsNull(add2_cnode) != lite::RET_OK || CheckInputSize(add2_cnode, kAddInputsLength) != lite::RET_OK) {
return nullptr;
}
auto add2_primitivec = GetValueNode<std::shared_ptr<lite::PrimitiveC>>(add2_cnode->input(0));
MS_ASSERT(utils::isa<std::shared_ptr<mindspore::lite::Add>>(add2_primitivec));
auto add2_op = utils::cast<std::shared_ptr<mindspore::lite::Add>>(add2_primitivec);
MS_ASSERT(add2_op != nullptr);
AnfNodePtr sub1_node = add2_cnode->input(2);
if (CheckIfAnfNodeIsNull(sub1_node) != lite::RET_OK) {
return nullptr;
}

// sub1
auto sub1_cnode = sub1_node->cast<CNodePtr>();
if (CheckIfCNodeIsNull(sub1_cnode) != lite::RET_OK || CheckInputSize(sub1_cnode, kSubInputsLength) != lite::RET_OK) {
return nullptr;
bool LayerNormFusion::GetAxis(const CNodePtr &input_cnode, const std::vector<int> &mean_axes,
const std::vector<int> &params_shape, int *begin_norm_axis,
int *begin_params_axis) const {
MS_ASSERT(input_node != nullptr);
MS_ASSERT(begin_norm_axis != nullptr);
MS_ASSERT(begin_params_axis != nullptr);
auto abstract = input_cnode->abstract();
if (abstract == nullptr) {
MS_LOG(DEBUG) << "abstract of input is nullptr";
return false;
}
if (!utils::isa<abstract::AbstractTensorPtr>(abstract)) {
MS_LOG(DEBUG) << "Abstract should be abstract tensor";
return false;
}
auto abstract_tensor = utils::cast<abstract::AbstractTensorPtr>(abstract);
if (!utils::isa<abstract::ShapePtr>(abstract_tensor->BuildShape())) {
MS_LOG(DEBUG) << "Shape of Abstract should be ShapePtr";
return false;
}
auto shape = utils::cast<abstract::ShapePtr>(abstract_tensor->BuildShape())->shape();
if (mean_axes.back() + 1 != static_cast<int>(shape.size())) {
MS_LOG(DEBUG) << "mean node is not reduce to last axis";
return false;
}
for (size_t i = 1; i < mean_axes.size(); ++i) {
if (mean_axes[i] != mean_axes[i - 1] + 1) {
MS_LOG(DEBUG) << "mean axes is not continuous";
return false;
}
}
auto sub1_primitivec = GetValueNode<std::shared_ptr<lite::PrimitiveC>>(sub1_cnode->input(0));
MS_ASSERT(utils::isa<std::shared_ptr<mindspore::lite::Sub>>(sub1_primitivec));
auto sub1_op = utils::cast<std::shared_ptr<mindspore::lite::Sub>>(sub1_primitivec);
MS_ASSERT(sub1_op != nullptr);
AnfNodePtr beta_node = sub1_cnode->input(1);
AnfNodePtr mul3_node = sub1_cnode->input(2);
if (CheckIfAnfNodeIsNull(beta_node) != lite::RET_OK || CheckIfAnfNodeIsNull(mul3_node) != lite::RET_OK) {
return nullptr;
// there is no need to check params_shape
*begin_norm_axis = mean_axes.front();
*begin_params_axis = static_cast<int>(shape.size()) - static_cast<int>(params_shape.size());
if (*begin_params_axis < 0) {
MS_LOG(DEBUG) << "LayerNorm begin_params_axis illegal, not fuse";
return false;
}
return true;
}

bool LayerNormFusion::CheckPattern(const EquivPtr &equiv, float *epsilon, int *begin_norm_axis,
int *begin_params_axis) const {
MS_ASSERT(equiv != nullptr);
MS_ASSERT(epsilon != nullptr);
MS_ASSERT(begin_norm_axis != nullptr);
MS_ASSERT(begin_params_axis != nullptr);
// beta
auto beta_node = utils::cast<AnfNodePtr>((*equiv)[beta_]);
MS_ASSERT(beta_node != nullptr);
if (CheckIfNodeIsParam(beta_node) != lite::RET_OK) {
return nullptr;
return false;
}
auto beta_param = beta_node->cast<ParameterPtr>()->default_param();
auto beta_tensor = std::dynamic_pointer_cast<ParamValueLite>(beta_param);
auto beta_shape = beta_tensor->tensor_shape();

// mul3
auto mul3_cnode = mul3_node->cast<CNodePtr>();
if (CheckIfCNodeIsNull(mul3_cnode) != lite::RET_OK || CheckInputSize(mul3_cnode, kMulInputsLength) != lite::RET_OK) {
return nullptr;
}
auto mul3_primitivec = GetValueNode<std::shared_ptr<lite::PrimitiveC>>(mul3_cnode->input(0));
MS_ASSERT(utils::isa<std::shared_ptr<mindspore::lite::Mul>>(mul3_primitivec));
auto mul3_op = utils::cast<std::shared_ptr<mindspore::lite::Mul>>(mul3_primitivec);
MS_ASSERT(mul3_op != nullptr);
AnfNodePtr mean1_node = mul3_cnode->input(1);
AnfNodePtr mul2_node = mul3_cnode->input(2);
if (CheckIfAnfNodeIsNull(mean1_node) != lite::RET_OK || CheckIfAnfNodeIsNull(mul2_node) != lite::RET_OK) {
return nullptr;
}

// mul2
auto mul2_cnode = mul2_node->cast<CNodePtr>();
if (CheckIfCNodeIsNull(mul2_cnode) != lite::RET_OK || CheckInputSize(mul2_cnode, kMulInputsLength) != lite::RET_OK) {
return nullptr;
}
auto mul2_primitivec = GetValueNode<std::shared_ptr<lite::PrimitiveC>>(mul2_cnode->input(0));
MS_ASSERT(utils::isa<std::shared_ptr<mindspore::lite::Mul>>(mul2_primitivec));
auto mul2_op = utils::cast<std::shared_ptr<mindspore::lite::Mul>>(mul2_primitivec);
MS_ASSERT(mul2_op != nullptr);
AnfNodePtr rsqrt_node = mul2_cnode->input(1);
AnfNodePtr gamma_node = mul2_cnode->input(2);
if (CheckIfAnfNodeIsNull(rsqrt_node) != lite::RET_OK || CheckIfAnfNodeIsNull(gamma_node) != lite::RET_OK) {
return nullptr;
}

// gamma
auto gamma_node = utils::cast<AnfNodePtr>((*equiv)[gamma_]);
MS_ASSERT(gamma_node != nullptr);
if (CheckIfNodeIsParam(gamma_node) != lite::RET_OK) {
return nullptr;
return false;
}
auto gamma_param = gamma_node->cast<ParameterPtr>()->default_param();
auto gamma_tensor = std::dynamic_pointer_cast<ParamValueLite>(gamma_param);
auto gamma_shape = gamma_tensor->tensor_shape();

// rsqrt
auto rsqrt_cnode = rsqrt_node->cast<CNodePtr>();
if (CheckIfCNodeIsNull(rsqrt_cnode) != lite::RET_OK ||
CheckInputSize(rsqrt_cnode, kRsqrtInputsLength) != lite::RET_OK) {
return nullptr;
}
auto rsqrt_primitivec = GetValueNode<std::shared_ptr<lite::PrimitiveC>>(rsqrt_cnode->input(0));
MS_ASSERT(utils::isa<std::shared_ptr<mindspore::lite::Rsqrt>>(rsqrt_primitivec));
auto rsqrt_op = utils::cast<std::shared_ptr<mindspore::lite::Rsqrt>>(rsqrt_primitivec);
MS_ASSERT(rsqrt_op != nullptr);
AnfNodePtr add1_node = rsqrt_cnode->input(1);
if (CheckIfAnfNodeIsNull(add1_node) != lite::RET_OK) {
return nullptr;
}

// add1
auto add1_cnode = add1_node->cast<CNodePtr>();
if (CheckIfCNodeIsNull(add1_cnode) != lite::RET_OK || CheckInputSize(add1_cnode, kAddInputsLength) != lite::RET_OK) {
return nullptr;
}
auto add1_primitivec = GetValueNode<std::shared_ptr<lite::PrimitiveC>>(add1_cnode->input(0));
MS_ASSERT(utils::isa<std::shared_ptr<mindspore::lite::Add>>(add1_primitivec));
auto add1_op = utils::cast<std::shared_ptr<mindspore::lite::Add>>(add1_primitivec);
MS_ASSERT(add1_op != nullptr);
AnfNodePtr mean2_node = add1_cnode->input(1);
AnfNodePtr epsilon_node = add1_cnode->input(2);
if (CheckIfAnfNodeIsNull(mean2_node) != lite::RET_OK || CheckIfAnfNodeIsNull(epsilon_node) != lite::RET_OK) {
return nullptr;
}

// epsilon
auto epsilon_node = utils::cast<AnfNodePtr>((*equiv)[epsilon_]);
MS_ASSERT(epsilon_node != nullptr);
if (CheckIfNodeIsParam(epsilon_node) != lite::RET_OK) {
// delete[] add_bias_data;
return nullptr;
return false;
}
auto epsilon_param = epsilon_node->cast<ParameterPtr>()->default_param();
auto epsilon_tensor = std::dynamic_pointer_cast<ParamValueLite>(epsilon_param);
auto epsilon_data = reinterpret_cast<float *>(epsilon_tensor->tensor_addr());
auto epsilon_shape = epsilon_tensor->tensor_shape();

// mean2
auto mean2_cnode = mean2_node->cast<CNodePtr>();
if (CheckIfCNodeIsNull(mean2_cnode) != lite::RET_OK ||
CheckInputSize(mean2_cnode, kReduceInputsLength) != lite::RET_OK) {
return nullptr;
auto mean2_value = utils::cast<AnfNodePtr>((*equiv)[mean2_]);
MS_ASSERT(mean2_value != nullptr);
auto mean2_primitivec = GetValueNode<std::shared_ptr<lite::PrimitiveC>>(mean2_value);
if (!utils::isa<std::shared_ptr<mindspore::lite::Reduce>>(mean2_primitivec)) {
return false;
}
auto mean2_primitivec = GetValueNode<std::shared_ptr<lite::PrimitiveC>>(mean2_cnode->input(0));
MS_ASSERT(utils::isa<std::shared_ptr<mindspore::lite::Reduce>>(mean2_primitivec));
auto mean2_op = utils::cast<std::shared_ptr<mindspore::lite::Reduce>>(mean2_primitivec);
MS_ASSERT(mean2_op != nullptr);
if (mean2_op->GetMode() != schema::ReduceMode_ReduceMean) {
return nullptr;
return false;
}
auto mean2_axes = mean2_op->GetAxes();
AnfNodePtr squared_difference_node = mean2_cnode->input(1);
if (CheckIfAnfNodeIsNull(squared_difference_node) != lite::RET_OK) {
return nullptr;
}

// mean1
auto mean1_cnode = mean1_node->cast<CNodePtr>();
if (CheckIfCNodeIsNull(mean1_cnode) != lite::RET_OK ||
CheckInputSize(mean1_cnode, kReduceInputsLength) != lite::RET_OK) {
return nullptr;
auto mean1_value = utils::cast<AnfNodePtr>((*equiv)[mean1_]);
MS_ASSERT(mean1_value != nullptr);
auto mean1_primitivec = GetValueNode<std::shared_ptr<lite::PrimitiveC>>(mean1_value);
if (!utils::isa<std::shared_ptr<mindspore::lite::Reduce>>(mean1_primitivec)) {
return false;
}
auto mean1_primitivec = GetValueNode<std::shared_ptr<lite::PrimitiveC>>(mean1_cnode->input(0));
MS_ASSERT(utils::isa<std::shared_ptr<mindspore::lite::Reduce>>(mean1_primitivec));
auto mean1_op = utils::cast<std::shared_ptr<mindspore::lite::Reduce>>(mean1_primitivec);
MS_ASSERT(mean1_op != nullptr);
if (mean1_op->GetMode() != schema::ReduceMode_ReduceMean) {
return nullptr;
return false;
}
AnfNodePtr input3_node = mean1_cnode->input(1);
auto mean1_axes = mean1_op->GetAxes();
if (CheckIfAnfNodeIsNull(input3_node) != lite::RET_OK) {
return nullptr;
}

// verify two mean ops have same axes
if (mean1_axes.size() != mean2_axes.size()) {
return nullptr;
auto input_node = utils::cast<AnfNodePtr>((*equiv)[input_]);
MS_ASSERT(input_node != nullptr);
if (!utils::isa<CNodePtr>(input_node)) {
return false;
}
for (size_t i = 0; i < mean1_axes.size(); ++i) {
if (mean1_axes[i] != mean2_axes[i]) {
return nullptr;
}
auto input_cnode = input_node->cast<CNodePtr>();
if (mean1_axes != mean2_axes) {
return false;
}
// verify axes size and gamma/beta size are equal
if (mean1_axes.size() != gamma_shape.size() || mean1_axes.size() != beta_shape.size()) {
return nullptr;
return false;
}
// verify gamma and beta have same shape
for (size_t i = 0; i < gamma_shape.size(); ++i) {
if (gamma_shape[i] != beta_shape[i]) {
return nullptr;
}
if (gamma_shape != beta_shape) {
return false;
}
// verify epsilon has exactly one element
float epsilon;
if (epsilon_shape.empty() || (epsilon_shape.size() == 1 && epsilon_shape[0] == 1)) {
epsilon = epsilon_data[0];
*epsilon = epsilon_data[0];
} else {
return nullptr;
return false;
}
if (!GetAxis(input_cnode, mean1_axes, gamma_shape, begin_norm_axis, begin_params_axis)) {
return false;
}
return true;
}

auto layer_norm_cnode = CreateLayerNormNode(func_graph, equiv, gamma_shape, epsilon);
const AnfNodePtr LayerNormFusion::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
const EquivPtr &equiv) const {
MS_ASSERT(func_graph != nullptr);
MS_ASSERT(node != nullptr);
MS_ASSERT(equiv != nullptr);
MS_LOG(DEBUG) << "layer_norm_fusion pass";
if (!utils::isa<CNodePtr>(node)) {
return nullptr;
}
auto add2_cnode = node->cast<CNodePtr>();
float epsilon = 0.0f;
int begin_norm_axis = 0;
int begin_params_axis = 0;
if (!CheckPattern(equiv, &epsilon, &begin_norm_axis, &begin_params_axis)) {
return nullptr;
}
auto layer_norm_cnode = CreateLayerNormNode(func_graph, equiv, epsilon, begin_norm_axis, begin_params_axis);
layer_norm_cnode->set_abstract(add2_cnode->abstract()->Clone());
layer_norm_cnode->set_fullname_with_scope("layer_norm_" + add2_cnode->fullname_with_scope());
MS_LOG(INFO) << "layernorm node:" << layer_norm_cnode->fullname_with_scope() << " fusion success";


+ 13
- 6
mindspore/lite/tools/optimizer/fusion/layer_norm_fusion.h View File

@@ -31,6 +31,8 @@ class LayerNormFusion : public PatternProcessPass {
explicit LayerNormFusion(const std::string &name = "layer_norm_fusion", bool multigraph = true)
: PatternProcessPass(name, multigraph) {
input_ = std::make_shared<Var>();
mean1_ = std::make_shared<Var>();
mean2_ = std::make_shared<Var>();
gamma_ = std::make_shared<Var>();
beta_ = std::make_shared<Var>();
epsilon_ = std::make_shared<Var>();
@@ -41,12 +43,17 @@ class LayerNormFusion : public PatternProcessPass {
const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override;

private:
CNodePtr CreateLayerNormNode(const FuncGraphPtr &func_graph, const EquivPtr &equiv, const std::vector<int> &shape,
const float epsilon) const;
VarPtr input_;
VarPtr gamma_;
VarPtr beta_;
VarPtr epsilon_;
bool GetAxis(const CNodePtr &input_cnode, const std::vector<int> &mean_axes, const std::vector<int> &params_shape,
int *begin_norm_axis, int *begin_params_axis) const;
bool CheckPattern(const EquivPtr &equiv, float *epsilon, int *begin_norm_axis, int *begin_params_axis) const;
CNodePtr CreateLayerNormNode(const FuncGraphPtr &func_graph, const EquivPtr &equiv, float epsilon,
int begin_norm_axis, int begin_params_axis) const;
VarPtr input_ = nullptr;
VarPtr mean1_ = nullptr;
VarPtr mean2_ = nullptr;
VarPtr gamma_ = nullptr;
VarPtr beta_ = nullptr;
VarPtr epsilon_ = nullptr;
};
} // namespace opt
} // namespace mindspore


Write Preview
Loading…
Cancel
Save