Browse Source

fix momentum fusion pass

tags/v1.1.0
wilfChen 5 years ago
parent
commit
3b7e01c698
4 changed files with 27 additions and 48 deletions
  1. +1
    -2
      mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_scale_fusion.cc
  2. +2
    -21
      mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_scale_fusion.h
  3. +22
    -0
      mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_weight_scale_fusion.cc
  4. +2
    -25
      mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_weight_scale_fusion.h

+ 1
- 2
mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_scale_fusion.cc View File

@@ -26,7 +26,6 @@

namespace mindspore {
namespace opt {
namespace {
bool ApplyMomentumScaleFusion::IsScalar(const BaseRef &n) {
if (utils::isa<AnfNodePtr>(n)) {
AnfNodePtr in = utils::cast<AnfNodePtr>(n);
@@ -86,5 +85,5 @@ const AnfNodePtr ApplyMomentumScaleFusion::Process(const FuncGraphPtr &graph, co
replace_node->set_scope(node->scope());
return replace_node;
}
} // namespace
} // namespace opt
} // namespace mindspore

+ 2
- 21
mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_scale_fusion.h View File

@@ -37,27 +37,8 @@ class ApplyMomentumScaleFusion : public PatternProcessPass {
const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override;

private:
static bool IsScalar(const BaseRef &n) {
if (utils::isa<AnfNodePtr>(n)) {
AnfNodePtr in = utils::cast<AnfNodePtr>(n);
MS_EXCEPTION_IF_NULL(in);
auto shape = in->Shape()->cast<abstract::ShapePtr>();
MS_EXCEPTION_IF_NULL(shape);
if (shape->shape().size() != 0) {
return false;
}
auto dtype = in->Type();
if (dtype->type_id() != kObjectTypeTensorType) {
return false;
}
auto element_type = dyn_cast<TensorType>(dtype)->element()->type_id();
if (element_type != kNumberTypeFloat32) {
return false;
}
return true;
}
return false;
}
static bool IsScalar(const BaseRef &n);

VarPtr scale_;
VarPtr variable_;
VarPtr accumulation_;


+ 22
- 0
mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_weight_scale_fusion.cc View File

@@ -26,6 +26,28 @@

namespace mindspore {
namespace opt {
bool ApplyMomentumWeightDecayScaleFusion::IsScalar(const BaseRef &n) {
if (utils::isa<AnfNodePtr>(n)) {
AnfNodePtr in = utils::cast<AnfNodePtr>(n);
MS_EXCEPTION_IF_NULL(in);
auto shape = in->Shape()->cast<abstract::ShapePtr>();
MS_EXCEPTION_IF_NULL(shape);
if (shape->shape().size() != 0) {
return false;
}
auto dtype = in->Type();
if (dtype->type_id() != kObjectTypeTensorType) {
return false;
}
auto element_type = dyn_cast<TensorType>(dtype)->element()->type_id();
if (element_type != kNumberTypeFloat32) {
return false;
}
return true;
}
return false;
}

const BaseRef ApplyMomentumWeightDecayScaleFusion::DefinePattern() const {
VectorRef weight = VectorRef(
{prim::kPrimAddN, VectorRef({prim::kPrimMul, variable_, weight_decay_}), VectorRef({prim::kPrimCast, gradient_})});


+ 2
- 25
mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_weight_scale_fusion.h View File

@@ -21,30 +21,6 @@

namespace mindspore {
namespace opt {
namespace {
bool IsScalar(const BaseRef &n) {
if (utils::isa<AnfNodePtr>(n)) {
AnfNodePtr in = utils::cast<AnfNodePtr>(n);
MS_EXCEPTION_IF_NULL(in);
auto shape = in->Shape()->cast<abstract::ShapePtr>();
MS_EXCEPTION_IF_NULL(shape);
if (shape->shape().size() != 0) {
return false;
}
auto dtype = in->Type();
if (dtype->type_id() != kObjectTypeTensorType) {
return false;
}
auto element_type = dyn_cast<TensorType>(dtype)->element()->type_id();
if (element_type != kNumberTypeFloat32) {
return false;
}
return true;
}
return false;
}
} // namespace

class ApplyMomentumWeightDecayScaleFusion : public PatternProcessPass {
public:
explicit ApplyMomentumWeightDecayScaleFusion(bool multigraph = true)
@@ -62,9 +38,10 @@ class ApplyMomentumWeightDecayScaleFusion : public PatternProcessPass {
const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override;

private:
static bool IsScalar(const BaseRef &n);

VarPtr weight_decay_;
VarPtr scale_;

VarPtr variable_;
VarPtr accumulation_;
VarPtr learning_rate_;


Loading…
Cancel
Save