Publishing 2019 R1 content

[platform/upstream/dldt.git] / model-optimizer / extensions / front / mxnet / check_softmax_node_inputs.py

diff --git a/model-optimizer/extensions/front/mxnet/check_softmax_node_inputs.py b/model-optimizer/extensions/front/mxnet/check_softmax_node_inputs.py
index 1e740ad..a8b6fa0 100644 (file)
--- a/model-optimizer/extensions/front/mxnet/check_softmax_node_inputs.py
+++ b/model-optimizer/extensions/front/mxnet/check_softmax_node_inputs.py
@@ -1,5 +1,5 @@
 """
- Copyright (c) 2017-2018 Intel Corporation
+ Copyright (c) 2017-2019 Intel Corporation
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -14,34 +14,39 @@
  limitations under the License.
 """
 
-import networkx as nx
 from mo.front.common.replacement import FrontReplacementPattern
+from mo.graph.graph import Graph
 
 
 class CheckSoftmaxNodeInputs(FrontReplacementPattern):
-
     enabled = True
 
+    def run_before(self):
+        from extensions.front.user_data_repack import UserDataRepack
+        return [UserDataRepack]
+
+    def run_after(self):
+        return []
+
     @staticmethod
     def pattern():
         return dict(
             nodes=[
-                ('softmax', dict(op='SoftmaxOutput'))
+                ('softmax', dict(op=lambda op: op in ['SoftMax', 'SoftmaxActivation', 'SoftmaxOutput']))
             ],
             edges=[])
 
     @staticmethod
-    def replace_pattern(graph: nx.MultiDiGraph, match: dict):
+    def replace_pattern(graph: Graph, match: dict):
         """
         Need to remove from softmax layer all unused inputs
         Parameters
         ----------
-        graph : nx.MultiDiGraph
+        graph : Graph
            Graph with loaded model.
          match : dict
            Patterns which were found in graph structure.
         """
-
         softmax_node = match['softmax']
         softmax_nodes_len = len(softmax_node.in_nodes())
         for i in reversed(range(1, softmax_nodes_len)):