From eb973172d29b23d131293b4cbde2e113b2914bd2 Mon Sep 17 00:00:00 2001 From: hwjiaorui Date: Mon, 4 Jan 2021 14:37:37 +0800 Subject: [PATCH] unify mindir sparse sparse loss format add string head context --- ..._cross_entropy_with_logits_unify_mindir.cc | 193 +++++++++++++++--- ...x_cross_entropy_with_logits_unify_mindir.h | 24 ++- .../ccsrc/backend/session/ascend_session.cc | 5 + mindspore/nn/loss/loss.py | 10 +- 4 files changed, 191 insertions(+), 41 deletions(-) diff --git a/mindspore/ccsrc/backend/optimizer/ascend/mindir/sparse_softmax_cross_entropy_with_logits_unify_mindir.cc b/mindspore/ccsrc/backend/optimizer/ascend/mindir/sparse_softmax_cross_entropy_with_logits_unify_mindir.cc index df9d59fbebb..7f79aed72b8 100644 --- a/mindspore/ccsrc/backend/optimizer/ascend/mindir/sparse_softmax_cross_entropy_with_logits_unify_mindir.cc +++ b/mindspore/ccsrc/backend/optimizer/ascend/mindir/sparse_softmax_cross_entropy_with_logits_unify_mindir.cc @@ -27,6 +27,9 @@ #include "ir/dtype/type.h" constexpr auto softmax_output_shape_size = 2; +constexpr auto kAttrDepth = "depth"; +constexpr auto kAttrMultiples = "multiples"; +constexpr auto kIsFeatureMapInputList = "IsFeatureMapInputList"; namespace mindspore { namespace opt { namespace { @@ -47,12 +50,12 @@ ValueNodePtr CreateValueNode(const ValuePtr &value_ptr, TypeId output_type) { return new_node; } -CNodePtr CreateOneHot(const FuncGraphPtr &graph, const CNodePtr &sparse_softmax_node) { +CNodePtr CreateOneHot(const FuncGraphPtr &graph, const CNodePtr &sparse_softmax_node, bool is_pynative = false) { MS_EXCEPTION_IF_NULL(graph); MS_EXCEPTION_IF_NULL(sparse_softmax_node); std::vector logits_shape = AnfAlgo::GetPrevNodeOutputInferShape(sparse_softmax_node, 0); - int64_t depth; + int64_t depth = 0; if (logits_shape.size() >= 1) { size_t index = logits_shape.size() - 1; depth = logits_shape[index]; @@ -66,33 +69,37 @@ CNodePtr CreateOneHot(const FuncGraphPtr &graph, const CNodePtr &sparse_softmax_ auto value_off = std::make_shared(0.0, kFloat32); auto value_off_node = CreateValueNode(value_off, kNumberTypeFloat32); MS_EXCEPTION_IF_NULL(value_off_node); - auto kernel_graph = graph->cast(); kernel_graph->AddValueNodeToGraph(value_on_node); kernel_graph->AddValueNodeToGraph(value_off_node); - auto depth_node = NewValueNode(depth); - MS_EXCEPTION_IF_NULL(depth_node); - - auto depth_abstract = std::make_shared(); - depth_abstract->set_type(kInt64); - depth_node->set_abstract(depth_abstract); - auto one_hot_primitive = std::make_shared(kOneHotOpName); std::vector input_names = {"indices", "depth", "on_value", "off_value"}; std::vector output_names = {"output"}; one_hot_primitive->set_attr(kAttrInputNames, MakeValue(input_names)); one_hot_primitive->set_attr(kAttrOutputNames, MakeValue(output_names)); - std::vector one_hot_inputs = {NewValueNode(one_hot_primitive), sparse_softmax_node->input(2), depth_node, - value_on_node, value_off_node}; + + std::vector one_hot_inputs; + if (is_pynative) { + one_hot_inputs = {NewValueNode(one_hot_primitive), sparse_softmax_node->input(2), value_on_node, value_off_node}; + } else { + auto depth_node = NewValueNode(depth); + MS_EXCEPTION_IF_NULL(depth_node); + auto depth_abstract = std::make_shared(); + depth_abstract->set_type(kInt64); + depth_node->set_abstract(depth_abstract); + one_hot_inputs = {NewValueNode(one_hot_primitive), sparse_softmax_node->input(2), depth_node, value_on_node, + value_off_node}; + } auto one_hot_node = graph->NewCNode(one_hot_inputs); MS_EXCEPTION_IF_NULL(one_hot_node); - one_hot_node->set_scope(sparse_softmax_node->scope()); std::vector labels_shape = AnfAlgo ::GetPrevNodeOutputInferShape(sparse_softmax_node, 1); labels_shape.emplace_back(depth); AnfAlgo::SetOutputInferTypeAndShape({kNumberTypeFloat32}, {labels_shape}, one_hot_node.get()); - AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue(-1), one_hot_node); + if (is_pynative) { + AnfAlgo::SetNodeAttr(kAttrDepth, MakeValue(depth), one_hot_node); + } return one_hot_node; } @@ -106,9 +113,6 @@ CNodePtr CreateSoftmaxCrossEntropyWithLogits(const FuncGraphPtr &graph, const CN MS_LOG(EXCEPTION) << "sparse_softmax_cross_entropy_with_logits's input size not equal " << kSparseSoftmaxCrossEntropyWithLogitsInputNum; } - if (one_hot_node->size() != kOneHotInputNum) { - MS_LOG(EXCEPTION) << "ont_hot's input size not equal " << kOneHotInputNum; - } std::vector inputs = {NewValueNode(std::make_shared(kSoftmaxCrossEntropyWithLogitsOpName)), sparse_softmax_node->input(1), one_hot_node}; @@ -131,7 +135,7 @@ CNodePtr CreateSoftmaxCrossEntropyWithLogits(const FuncGraphPtr &graph, const CN return softmax_node; } -ValueNodePtr GetAxis(const AnfNodePtr &node) { +std::vector GetAxis(const AnfNodePtr &node) { MS_EXCEPTION_IF_NULL(node); std::vector output_shape = AnfAlgo::GetOutputInferShape(node, 0); if (output_shape.empty()) { @@ -141,13 +145,19 @@ ValueNodePtr GetAxis(const AnfNodePtr &node) { for (size_t i = 0; i < output_shape.size(); i++) { range.emplace_back(i); } + return range; +} + +ValueNodePtr GetAxisNode(const AnfNodePtr &node) { + MS_EXCEPTION_IF_NULL(node); + auto range = GetAxis(node); auto axis_node = CreateValueNode(MakeValue(range), kNumberTypeInt64); MS_EXCEPTION_IF_NULL(axis_node); return axis_node; } CNodePtr CreateReduceMean(const FuncGraphPtr &graph, const CNodePtr &sparse_softmax_node, - const AnfNodePtr &softmax_output_node) { + const AnfNodePtr &softmax_output_node, bool is_pynative = false) { MS_EXCEPTION_IF_NULL(graph); MS_EXCEPTION_IF_NULL(sparse_softmax_node); MS_EXCEPTION_IF_NULL(softmax_output_node); @@ -155,10 +165,10 @@ CNodePtr CreateReduceMean(const FuncGraphPtr &graph, const CNodePtr &sparse_soft MS_LOG(EXCEPTION) << "sparse_softmax_cross_entropy_with_logits's input size not equal " << kSparseSoftmaxCrossEntropyWithLogitsInputNum; } - auto axis_node = GetAxis(softmax_output_node); + + auto axis_value = GetAxis(softmax_output_node); + auto axis_node = GetAxisNode(softmax_output_node); MS_EXCEPTION_IF_NULL(axis_node); - auto kernel_graph = graph->cast(); - kernel_graph->AddValueNodeToGraph(axis_node); auto reduce_primitive = std::make_shared(kReduceMeanOpName); std::vector input_names = {"x", "axis"}; @@ -166,14 +176,23 @@ CNodePtr CreateReduceMean(const FuncGraphPtr &graph, const CNodePtr &sparse_soft reduce_primitive->set_attr(kAttrInputNames, MakeValue(input_names)); reduce_primitive->set_attr(kAttrOutputNames, MakeValue(output_names)); - std::vector inputs = {NewValueNode(reduce_primitive), softmax_output_node, axis_node}; + auto kernel_graph = graph->cast(); + std::vector inputs; + if (is_pynative) { + inputs = {NewValueNode(reduce_primitive), softmax_output_node}; + } else { + kernel_graph->AddValueNodeToGraph(axis_node); + inputs = {NewValueNode(reduce_primitive), softmax_output_node, axis_node}; + } auto reduce_node = graph->NewCNode(inputs); MS_EXCEPTION_IF_NULL(reduce_node); - reduce_node->set_scope(sparse_softmax_node->scope()); auto reduce_abstract = softmax_output_node->abstract(); reduce_abstract->set_shape(std::make_shared()); reduce_node->set_abstract(reduce_abstract); + if (is_pynative) { + AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue(axis_value), reduce_node); + } return reduce_node; } @@ -207,8 +226,33 @@ CNodePtr CreateExpandDims(const FuncGraphPtr &graph, const CNodePtr &real_div_no expand_dims_node.get()); return expand_dims_node; } +CNodePtr CreateExpandDimsPynative(const FuncGraphPtr &graph, const CNodePtr &real_div_node) { + MS_EXCEPTION_IF_NULL(graph); + MS_EXCEPTION_IF_NULL(real_div_node); + if (real_div_node->size() != kRealDivInputNum) { + MS_LOG(EXCEPTION) << "Op real_div's input num not equal " << kRealDivInputNum; + } + int64_t axis = -1; + auto expand_dims_primitive = std::make_shared(kExpandDimsOpName); + std::vector input_names = {"x"}; + std::vector output_names = {"output"}; + expand_dims_primitive->set_attr(kAttrInputNames, MakeValue(input_names)); + expand_dims_primitive->set_attr(kAttrOutputNames, MakeValue(output_names)); + std::vector expand_dims_inputs = {NewValueNode(expand_dims_primitive), real_div_node}; + auto expand_dims_node = graph->NewCNode(expand_dims_inputs); + MS_EXCEPTION_IF_NULL(expand_dims_node); -CNodePtr CreateTile(const FuncGraphPtr &graph, const CNodePtr &sparse_softmax_node, const CNodePtr &mul_node) { + expand_dims_node->set_scope(real_div_node->scope()); + std::vector y_shape = AnfAlgo::GetOutputInferShape(real_div_node, 0); + y_shape.emplace_back(1); + AnfAlgo::SetOutputInferTypeAndShape({AnfAlgo::GetOutputInferDataType(real_div_node, 0)}, {y_shape}, + expand_dims_node.get()); + AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue(axis), expand_dims_node); + return expand_dims_node; +} + +CNodePtr CreateTile(const FuncGraphPtr &graph, const CNodePtr &sparse_softmax_node, const CNodePtr &mul_node, + bool is_pynative = false) { MS_EXCEPTION_IF_NULL(graph); MS_EXCEPTION_IF_NULL(sparse_softmax_node); MS_EXCEPTION_IF_NULL(mul_node); @@ -224,24 +268,37 @@ CNodePtr CreateTile(const FuncGraphPtr &graph, const CNodePtr &sparse_softmax_no std::vector multiple_value; std::transform(labels_shape.begin(), labels_shape.end(), std::back_inserter(multiple_value), [](size_t label) { return static_cast(label); }); - auto mutiples = MakeValue(multiple_value); - auto mutiples_node = CreateValueNode(mutiples, kNumberTypeInt64); - MS_EXCEPTION_IF_NULL(mutiples_node); - auto kernel_graph = graph->cast(); - kernel_graph->AddValueNodeToGraph(mutiples_node); + auto multiples = MakeValue(multiple_value); + auto multiples_node = CreateValueNode(multiples, kNumberTypeInt64); + MS_EXCEPTION_IF_NULL(multiples_node); auto tile_primitive = std::make_shared(kTileOpName); std::vector input_names = {"x", "multiples"}; std::vector output_names = {"output"}; tile_primitive->set_attr(kAttrInputNames, MakeValue(input_names)); tile_primitive->set_attr(kAttrOutputNames, MakeValue(output_names)); - std::vector tile_inputs = {NewValueNode(tile_primitive), mul_node->input(2), mutiples_node}; + + std::vector tile_inputs; + if (is_pynative) { + tile_inputs = {NewValueNode(tile_primitive), mul_node->input(2)}; + } else { + auto kernel_graph = graph->cast(); + kernel_graph->AddValueNodeToGraph(multiples_node); + tile_inputs = {NewValueNode(tile_primitive), mul_node->input(2), multiples_node}; + } + auto tile_node = graph->NewCNode(tile_inputs); MS_EXCEPTION_IF_NULL(tile_node); - tile_node->set_scope(mul_node->scope()); AnfAlgo::SetOutputInferTypeAndShape({AnfAlgo::GetPrevNodeOutputInferDataType(mul_node, 1)}, {labels_shape}, tile_node.get()); + if (is_pynative) { + AnfAlgo::SetNodeAttr(kAttrMultiples, MakeValue(multiples), tile_node); + } + // feature map set + std::vector feature_map_input_indexs; + feature_map_input_indexs.push_back(0); + AnfAlgo::SetNodeAttr(kIsFeatureMapInputList, MakeValue(feature_map_input_indexs), tile_node); return tile_node; } @@ -368,7 +425,6 @@ const AnfNodePtr SparseSoftmaxCrossEntropyWithLogitsUnifyMindIR::Process(const F std::vector softmax_node_outputs; CreateMultipleOutputsOfAnfNode(graph, softmax_node, kSoftmaxCrossEntropyWithLogitsOutputNum, &softmax_node_outputs); auto reduce_node = CreateReduceMean(graph, sparse_softmax_node, softmax_node_outputs[0]); - return reduce_node; } @@ -450,5 +506,76 @@ const AnfNodePtr GradSparseSoftmaxCrossEntropyWithLogitsUnifyMindIRV2::Process(c manager->Replace(sparse_softmax_node_grad, softmax_node_outputs[1]); return mul_node; } + +const AnfNodePtr PynativeSparseSoftmaxCrossEntropyWithLogitsUnifyMindIR::Process(const FuncGraphPtr &graph, + const AnfNodePtr &node, + const EquivPtr &) const { + MS_EXCEPTION_IF_NULL(graph); + MS_EXCEPTION_IF_NULL(node); + + auto sparse_softmax_node = node->cast(); + MS_EXCEPTION_IF_NULL(sparse_softmax_node); + if (sparse_softmax_node->size() != kSparseSoftmaxCrossEntropyWithLogitsInputNum) { + MS_LOG(EXCEPTION) << "Op SparseSoftmaxCrossEntropyWithLogits's input not equal " + << kSparseSoftmaxCrossEntropyWithLogitsInputNum; + } + if (AnfAlgo::HasNodeAttr(kAttrIsGrad, sparse_softmax_node) && + AnfAlgo::GetNodeAttr(sparse_softmax_node, kAttrIsGrad)) { + return nullptr; + } + + CNodePtr softmax_node; + auto one_hot_node = CreateOneHot(graph, sparse_softmax_node, true); + softmax_node = CreateSoftmaxCrossEntropyWithLogits(graph, sparse_softmax_node, one_hot_node); + + std::vector softmax_node_outputs; + CreateMultipleOutputsOfAnfNode(graph, softmax_node, kSoftmaxCrossEntropyWithLogitsOutputNum, &softmax_node_outputs); + auto reduce_node = CreateReduceMean(graph, sparse_softmax_node, softmax_node_outputs[0], true); + return reduce_node; +} + +const BaseRef PynativeGradSparseSoftmaxCrossEntropyWithLogitsUnifyMindIR::DefinePattern() const { + VarPtr x1 = std::make_shared(); + VarPtr x2 = std::make_shared(); + VarPtr x3 = std::make_shared(); + VectorRef sparse_softmax_cross_entropy_with_logits({prim::kPrimSparseSoftmaxCrossEntropyWithLogits, x1, x2}); + return VectorRef({prim::kPrimMul, sparse_softmax_cross_entropy_with_logits, x3}); +} + +const AnfNodePtr PynativeGradSparseSoftmaxCrossEntropyWithLogitsUnifyMindIR::Process(const FuncGraphPtr &graph, + const AnfNodePtr &node, + const EquivPtr &) const { + MS_EXCEPTION_IF_NULL(graph); + MS_EXCEPTION_IF_NULL(node); + + auto mul_node = node->cast(); + MS_EXCEPTION_IF_NULL(mul_node); + if (mul_node->size() != kMulInputNum) { + MS_LOG(EXCEPTION) << "Op Mul's input not equal " << kMulInputNum; + } + auto sparse_softmax_node = mul_node->input(1); + auto sparse_softmax_node_grad = sparse_softmax_node->cast(); + MS_EXCEPTION_IF_NULL(sparse_softmax_node_grad); + + if (sparse_softmax_node_grad->size() != kSparseSoftmaxCrossEntropyWithLogitsInputNum) { + MS_LOG(EXCEPTION) << "Op SparseSoftmaxCrossEntropyWithLogits's input not equal " + << kSparseSoftmaxCrossEntropyWithLogitsInputNum; + } + + CNodePtr softmax_node; + auto one_hot_node = CreateOneHot(graph, sparse_softmax_node_grad, true); + softmax_node = CreateSoftmaxCrossEntropyWithLogits(graph, sparse_softmax_node_grad, one_hot_node); + + std::vector softmax_node_outputs; + CreateMultipleOutputsOfAnfNode(graph, softmax_node, kSoftmaxCrossEntropyWithLogitsOutputNum, &softmax_node_outputs); + auto tile_node = CreateTile(graph, sparse_softmax_node_grad, mul_node, true); + auto real_div_node = CreateRealDiv(graph, sparse_softmax_node_grad, tile_node); + auto expand_dims_node = CreateExpandDimsPynative(graph, real_div_node); + + mul_node->set_input(1, softmax_node_outputs[1]); + mul_node->set_input(2, expand_dims_node); + + return mul_node; +} } // namespace opt } // namespace mindspore diff --git a/mindspore/ccsrc/backend/optimizer/ascend/mindir/sparse_softmax_cross_entropy_with_logits_unify_mindir.h b/mindspore/ccsrc/backend/optimizer/ascend/mindir/sparse_softmax_cross_entropy_with_logits_unify_mindir.h index 27e2837ff9f..2a321a7a1fe 100644 --- a/mindspore/ccsrc/backend/optimizer/ascend/mindir/sparse_softmax_cross_entropy_with_logits_unify_mindir.h +++ b/mindspore/ccsrc/backend/optimizer/ascend/mindir/sparse_softmax_cross_entropy_with_logits_unify_mindir.h @@ -18,14 +18,16 @@ #define MINDSPORE_SPARSE_SOFTMAX_CROSS_ENTROPY_WITH_LOGITS_UNIFY_MINDIR_H #include +#include #include "backend/optimizer/common/optimizer.h" namespace mindspore { namespace opt { class SparseSoftmaxCrossEntropyWithLogitsUnifyMindIR : public PatternProcessPass { public: - explicit SparseSoftmaxCrossEntropyWithLogitsUnifyMindIR(bool multigraph = true) - : PatternProcessPass("sparse_softmax_cross_entropy_with_logits_unify_mindir", multigraph) {} + explicit SparseSoftmaxCrossEntropyWithLogitsUnifyMindIR( + const std::string &name = "sparse_softmax_cross_entropy_with_logits_unify_mindir", bool multigraph = true) + : PatternProcessPass(name, multigraph) {} ~SparseSoftmaxCrossEntropyWithLogitsUnifyMindIR() override = default; const BaseRef DefinePattern() const override; const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override; @@ -49,6 +51,24 @@ class GradSparseSoftmaxCrossEntropyWithLogitsUnifyMindIRV2 : public PatternProce const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override; }; +class PynativeSparseSoftmaxCrossEntropyWithLogitsUnifyMindIR : public SparseSoftmaxCrossEntropyWithLogitsUnifyMindIR { + public: + explicit PynativeSparseSoftmaxCrossEntropyWithLogitsUnifyMindIR(bool multigraph = true) + : SparseSoftmaxCrossEntropyWithLogitsUnifyMindIR("pynative_sparse_softmax_cross_entropy_with_logits_unify_mindir", + multigraph) {} + ~PynativeSparseSoftmaxCrossEntropyWithLogitsUnifyMindIR() override = default; + const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override; +}; + +class PynativeGradSparseSoftmaxCrossEntropyWithLogitsUnifyMindIR : public PatternProcessPass { + public: + explicit PynativeGradSparseSoftmaxCrossEntropyWithLogitsUnifyMindIR(bool multigraph = true) + : PatternProcessPass("pynative_grad_sparse_softmax_cross_entropy_with_logits_unify_mindir", multigraph) {} + ~PynativeGradSparseSoftmaxCrossEntropyWithLogitsUnifyMindIR() override = default; + const BaseRef DefinePattern() const override; + const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override; +}; + } // namespace opt } // namespace mindspore #endif // MINDSPORE_SPARSE_SOFTMAX_CROSS_ENTROPY_WITH_LOGITS_UNIFY_MINDIR_H diff --git a/mindspore/ccsrc/backend/session/ascend_session.cc b/mindspore/ccsrc/backend/session/ascend_session.cc index a110d45e63c..6712e2ad858 100644 --- a/mindspore/ccsrc/backend/session/ascend_session.cc +++ b/mindspore/ccsrc/backend/session/ascend_session.cc @@ -449,9 +449,14 @@ void AscendSession::UnifyMindIR(const KernelGraphPtr &graph) { if (ms_context->get_param(MS_CTX_EXECUTION_MODE) == kGraphMode) { unify_mindir_pm->AddPass(std::make_shared()); unify_mindir_pm->AddPass(std::make_shared()); + unify_mindir_pm->AddPass(std::make_shared()); + unify_mindir_pm->AddPass(std::make_shared()); + unify_mindir_pm->AddPass(std::make_shared()); } else { unify_mindir_pm->AddPass(std::make_shared()); unify_mindir_pm->AddPass(std::make_shared()); + unify_mindir_pm->AddPass(std::make_shared()); + unify_mindir_pm->AddPass(std::make_shared()); } optimizer->AddPassManager(unify_mindir_pm); diff --git a/mindspore/nn/loss/loss.py b/mindspore/nn/loss/loss.py index 16f3dcb1744..3f32d4ca711 100644 --- a/mindspore/nn/loss/loss.py +++ b/mindspore/nn/loss/loss.py @@ -281,15 +281,13 @@ class SoftmaxCrossEntropyWithLogits(_Loss): self.on_value = Tensor(1.0, mstype.float32) self.off_value = Tensor(0., mstype.float32) self.is_cpugpu = context.get_context('device_target') in ["CPU", "GPU"] - if self.is_cpugpu: - self.sparse_softmax_cross_entropy = P.SparseSoftmaxCrossEntropyWithLogits() + self.sparse_softmax_cross_entropy = P.SparseSoftmaxCrossEntropyWithLogits() def construct(self, logits, labels): - if self.is_cpugpu and self.sparse and self.reduction == 'mean': - x = self.sparse_softmax_cross_entropy(logits, labels) - return x - if self.sparse: + if self.reduction == 'mean': + x = self.sparse_softmax_cross_entropy(logits, labels) + return x labels = self.one_hot(labels, F.shape(logits)[-1], self.on_value, self.off_value) x = self.softmax_cross_entropy(logits, labels)[0] return self.get_loss(x)