From 041f628ccef3588b888102a5b3ce940dc1d0f3f6 Mon Sep 17 00:00:00 2001 From: Kang Date: Tue, 2 Jun 2020 20:32:25 +0800 Subject: [PATCH] Optimize MixedPrecisionCast function --- mindspore/ccsrc/operator/ops.cc | 1 + mindspore/ccsrc/operator/ops.h | 1 + mindspore/ccsrc/pipeline/parse/parse.cc | 2 +- .../ccsrc/pipeline/static_analysis/prim.cc | 63 +++++++++++++++++-- .../ccsrc/pipeline/static_analysis/prim.h | 16 +++++ .../static_analysis/static_analysis.cc | 4 ++ mindspore/nn/wrap/cell_wrapper.py | 3 +- mindspore/ops/composite/base.py | 30 --------- mindspore/ops/functional.py | 1 + mindspore/train/amp.py | 3 +- 10 files changed, 83 insertions(+), 41 deletions(-) diff --git a/mindspore/ccsrc/operator/ops.cc b/mindspore/ccsrc/operator/ops.cc index 24083d36243..06d84dd4e3e 100755 --- a/mindspore/ccsrc/operator/ops.cc +++ b/mindspore/ccsrc/operator/ops.cc @@ -242,6 +242,7 @@ const PrimitivePtr kPrimIs_ = std::make_shared("is_"); const PrimitivePtr kPrimIsNot = std::make_shared("is_not"); const PrimitivePtr kPrimInDict = std::make_shared("in_dict"); const PrimitivePtr kPrimNotInDict = std::make_shared("not_in_dict"); +const PrimitivePtr kPrimMixedPrecisionCast = std::make_shared("mixed_precision_cast"); // Comm ops const PrimitivePtr kPrimMirror = std::make_shared("_MirrorOperator"); diff --git a/mindspore/ccsrc/operator/ops.h b/mindspore/ccsrc/operator/ops.h index b8c87edff40..e372a8ab1fd 100755 --- a/mindspore/ccsrc/operator/ops.h +++ b/mindspore/ccsrc/operator/ops.h @@ -251,6 +251,7 @@ extern const PrimitivePtr kPrimIs_; extern const PrimitivePtr kPrimIsNot; extern const PrimitivePtr kPrimInDict; extern const PrimitivePtr kPrimNotInDict; +extern const PrimitivePtr kPrimMixedPrecisionCast; // Comm ops extern const PrimitivePtr kPrimMirror; diff --git a/mindspore/ccsrc/pipeline/parse/parse.cc b/mindspore/ccsrc/pipeline/parse/parse.cc index c6e5d3713aa..dc7c5aa4a0b 100644 --- a/mindspore/ccsrc/pipeline/parse/parse.cc +++ b/mindspore/ccsrc/pipeline/parse/parse.cc @@ -67,7 +67,7 @@ AnfNodePtr GetMixedPrecisionCastHelp(const FuncGraphPtr &func_graph, const AnfNo } else { return param; } - auto cast_helper = prim::GetPythonOps("_mp_cast_helper", "mindspore.ops.composite.base"); + auto cast_helper = prim::kPrimMixedPrecisionCast; auto cast = func_graph->NewCNode({NewValueNode(cast_helper), NewValueNode(dst_type), param}); return cast; } diff --git a/mindspore/ccsrc/pipeline/static_analysis/prim.cc b/mindspore/ccsrc/pipeline/static_analysis/prim.cc index bcd02884240..a0bcaab5b7e 100644 --- a/mindspore/ccsrc/pipeline/static_analysis/prim.cc +++ b/mindspore/ccsrc/pipeline/static_analysis/prim.cc @@ -147,9 +147,6 @@ EvalResultPtr StandardPrimEvaluator::EvalPrim(const AnalysisEnginePtr &engine, c EvalResultPtr DoSignatureEvaluator::Run(AnalysisEnginePtr engine, const ConfigPtrList &args_conf_list, AnfNodeConfigPtr out_conf) { AbstractBasePtrList args_spec_list; - if (!prim_->isa()) { - MS_LOG(EXCEPTION) << "Primitive should be DoSignature, but " << prim_->ToString(); - } if (out_conf->node() == nullptr || !out_conf->node()->isa()) { MS_LOG(EXCEPTION) << "Node of out_conf should be CNode"; } @@ -221,9 +218,6 @@ EvalResultPtr UnpackGraphEvaluator::Run(AnalysisEnginePtr engine, const ConfigPt if (out_conf->node() == nullptr || !out_conf->node()->isa()) { MS_LOG(EXCEPTION) << "Node of out_conf should be CNode"; } - if (!prim_->isa()) { - MS_LOG(EXCEPTION) << "Primitive should be UnpackGraphPrimitive, but got " << prim_->ToString(); - } auto unpack_graph = prim_->cast(); auto out_node = out_conf->node()->cast(); @@ -267,6 +261,63 @@ EvalResultPtr UnpackGraphEvaluator::Run(AnalysisEnginePtr engine, const ConfigPt return engine->ForwardConfig(out_conf, fn_conf); } +AnfNodePtr MixedPrecisionCastHelper(AnfNodePtr source_node, AbstractBasePtr node_type, AnfNodePtr target_type, + FuncGraphPtr func_graph) { + AnfNodePtr target_node = source_node; + if (node_type->isa()) { + auto x = node_type->cast(); + if (x->element()->BuildType()->isa()) { + auto cast = prim::GetPythonOps("cast", "mindspore.ops.functional"); + MS_EXCEPTION_IF_NULL(cast); + target_node = func_graph->NewCNode({NewValueNode(cast), source_node, target_type}); + } + } else if (node_type->isa()) { + auto x = node_type->cast(); + auto &items = x->elements(); + std::size_t size = items.size(); + std::vector nodes; + nodes.emplace_back(NewValueNode(prim::kPrimMakeTuple)); + for (int i = 0; i < SizeToInt(size); i++) { + AnfNodePtr tuple_node = + func_graph->NewCNode({NewValueNode(prim::kPrimTupleGetItem), source_node, NewValueNode(i)}); + AnfNodePtr node = MixedPrecisionCastHelper(tuple_node, items[i], target_type, func_graph); + nodes.emplace_back(node); + } + target_node = func_graph->NewCNode(nodes); + } + return target_node; +} + +EvalResultPtr MixedPrecisionCastEvaluator::Run(AnalysisEnginePtr engine, const ConfigPtrList &args_conf_list, + AnfNodeConfigPtr out_conf) { + AbstractBasePtrList args_spec_list; + if (out_conf->node() == nullptr || !out_conf->node()->isa()) { + MS_LOG(EXCEPTION) << "Node of out_conf should be CNode"; + } + auto out_node = out_conf->node()->cast(); + const auto &out_node_inputs = out_node->inputs(); + if (out_node->inputs().size() == 0 || (out_node_inputs.size() - 1) != args_conf_list.size()) { + MS_LOG(EXCEPTION) << "MixedPrecisionCast" + << " args size should equal to inputs size minus 1, but args size " << args_conf_list.size() + << ", inputs size " << out_node_inputs.size(); + } + AnfNodePtrList args_inputs{out_node_inputs.begin() + 1, out_node_inputs.end()}; + (void)std::transform(args_conf_list.begin(), args_conf_list.end(), std::back_inserter(args_spec_list), + [](const ConfigPtr &ref) -> AbstractBasePtr { return ref->GetEvaluatedValue()->abstract(); }); + + ScopePtr scope = kDefaultScope; + if (out_conf != nullptr) { + scope = out_conf->node()->scope(); + } + ScopeGuard scope_guard(scope); + + FuncGraphPtr func_graph = out_conf->node()->func_graph(); + AnfNodePtr new_node = MixedPrecisionCastHelper(out_node_inputs[2], args_spec_list[1], out_node_inputs[1], func_graph); + AnfNodeConfigPtr fn_conf = engine->MakeConfig(new_node, out_conf->context()); + + return engine->ForwardConfig(out_conf, fn_conf); +} + namespace { py::object BuildValue(const ValuePtr &value_ptr) { if (value_ptr == nullptr) { diff --git a/mindspore/ccsrc/pipeline/static_analysis/prim.h b/mindspore/ccsrc/pipeline/static_analysis/prim.h index 219d9604b4b..3433b86cca6 100644 --- a/mindspore/ccsrc/pipeline/static_analysis/prim.h +++ b/mindspore/ccsrc/pipeline/static_analysis/prim.h @@ -102,6 +102,22 @@ class UnpackGraphEvaluator : public Evaluator { PrimitivePtr prim_; }; +class MixedPrecisionCastEvaluator : public Evaluator { + public: + explicit MixedPrecisionCastEvaluator(const PrimitivePtr primitive) + : Evaluator("MixedPrecisionCastEvaluator"), prim_(primitive) {} + ~MixedPrecisionCastEvaluator() override = default; + EvalResultPtr Run(AnalysisEnginePtr engine, const ConfigPtrList &argrefs, + AnfNodeConfigPtr out_config = nullptr) override; + + EvalResultPtr Eval(AnalysisEnginePtr, const AbstractBasePtrList &) override { + MS_LOG(EXCEPTION) << "Eval() should not be called, Run() method should be called"; + } + + private: + PrimitivePtr prim_; +}; + bool IsInWhiteList(PrimitivePtr primitive); StandardPrimitiveEvalImpl GetPrimitiveInferImpl(const PrimitivePtr &primitive); diff --git a/mindspore/ccsrc/pipeline/static_analysis/static_analysis.cc b/mindspore/ccsrc/pipeline/static_analysis/static_analysis.cc index b7520176ec6..48bcff62433 100644 --- a/mindspore/ccsrc/pipeline/static_analysis/static_analysis.cc +++ b/mindspore/ccsrc/pipeline/static_analysis/static_analysis.cc @@ -308,6 +308,10 @@ EvaluatorPtr GetPrimEvaluator(const PrimitivePtr &prim, const AnalysisEnginePtr evaluator = std::make_shared(prim); return evaluator; } + if (prim->name() == prim::kPrimMixedPrecisionCast->name()) { + evaluator = std::make_shared(prim); + return evaluator; + } if (prim->HasPyEvaluator()) { auto prim_py = dyn_cast(prim); if (prim_py != nullptr) { diff --git a/mindspore/nn/wrap/cell_wrapper.py b/mindspore/nn/wrap/cell_wrapper.py index a52f9e0d1f5..f0d920f51fa 100644 --- a/mindspore/nn/wrap/cell_wrapper.py +++ b/mindspore/nn/wrap/cell_wrapper.py @@ -21,7 +21,6 @@ from ...common.parameter import Parameter, ParameterTuple from ...ops import composite as C from ...ops import functional as F from ...ops import operations as P -from ...ops.composite.base import _mp_cast_helper from ...ops.operations.comm_ops import _VirtualDataset from ..cell import Cell from .grad_reducer import DistributedGradReducer @@ -345,7 +344,7 @@ class WithEvalCell(Cell): def construct(self, data, label): outputs = self._network(data) if self.add_cast_fp32: - label = _mp_cast_helper(mstype.float32, label) + label = F.mixed_precision_cast(mstype.float32, label) outputs = F.cast(outputs, mstype.float32) loss = self._loss_fn(outputs, label) return loss, outputs, label diff --git a/mindspore/ops/composite/base.py b/mindspore/ops/composite/base.py index 416ec33081f..007d147bf08 100644 --- a/mindspore/ops/composite/base.py +++ b/mindspore/ops/composite/base.py @@ -24,7 +24,6 @@ from ..._c_expression import EnvInstance_, GradOperation_, HyperMap_, MultitypeF from ...common import dtype as mstype from ...common.api import ms_function, _pynative_exec from .. import functional as F -from .. import operations as P from ...common.parameter import Parameter @@ -297,32 +296,3 @@ env_get = MultitypeFuncGraph("env_get") def _tensor_env_get(env, parameter): """Used to get env.""" return F.env_getitem(env, F.ref_to_embed(parameter), F.zeros_like(parameter)) - - -_mp_cast_helper = MultitypeFuncGraph('mixed_precision_cast_helper') - - -@_mp_cast_helper.register("TypeType", "Number") -@core -def _mixed_precision_cast_helper_1(type_, x): - """if x is float cast to type.""" - # type_ is place holder - return x - - -@_mp_cast_helper.register("TypeType", "Tensor") -@core -def _mixed_precision_cast_helper_2(type_, x): - """if x is float cast to type.""" - if F.issubclass_(F.dtype(x), mstype.float_): - return P.Cast()(x, type_) - return x - -@_mp_cast_helper.register("TypeType", "Tuple") -@core -def _mixed_precision_cast_helper_3(type_, x): - """if x is a tuple""" - t = () - for item in x: - t = t + (_mp_cast_helper(type_, item),) - return t diff --git a/mindspore/ops/functional.py b/mindspore/ops/functional.py index 8f5fcaefbc1..fa958a49e8b 100644 --- a/mindspore/ops/functional.py +++ b/mindspore/ops/functional.py @@ -126,6 +126,7 @@ is_ = Primitive("is_") is_not = Primitive("is_not") in_dict = Primitive("in_dict") not_in_dict = Primitive("not_in_dict") +mixed_precision_cast = Primitive("mixed_precision_cast") broadcast_gradient_args = Primitive('BroadcastGradientArgs') dot = Primitive('dot') array_reduce = Primitive('array_reduce') diff --git a/mindspore/train/amp.py b/mindspore/train/amp.py index da0626d6e8e..2baa1f591a3 100644 --- a/mindspore/train/amp.py +++ b/mindspore/train/amp.py @@ -21,7 +21,6 @@ from .._checkparam import Rel from ..common import dtype as mstype from ..nn.wrap.cell_wrapper import _VirtualDatasetCell from ..ops import functional as F -from ..ops.composite.base import _mp_cast_helper from ..parallel._utils import _get_parallel_mode from .loss_scale_manager import DynamicLossScaleManager, LossScaleManager from .parallel_utils import ParallelMode @@ -98,7 +97,7 @@ def _add_loss_network(network, loss_fn, cast_model_type): def construct(self, data, label): out = self._backbone(data) - label = _mp_cast_helper(mstype.float32, label) + label = F.mixed_precision_cast(mstype.float32, label) return self._loss_fn(F.cast(out, mstype.float32), label) validator.check_value_type('loss_fn', loss_fn, nn.Cell, None)