fix for different topology of gelu and layernorm fusion in Py2

Author: slin07

coremltools/converters/nnssa/coreml/graph_pass/op_fusions.py

@@ -240,6 +240,12 @@ def onehot_matmul_to_embedding(nnssa):
                 print('[Op Fusion] Node %s is removed.' %(inp_node.name))
 
 
+def _search_nodes_by_type(gf, node_names, op_type):
+    for name in node_names:
+        if gf[name].op == op_type:
+            return gf[name]
+
+
 def _match_layernorm_pattern(gf, entry_node):
     """ Return the nodes that form the subgraph of a LayerNormalization layer
     """
@@ -248,7 +254,10 @@ def _axes_in_range(axes, rank):
 
     try:
         params = {}
-        mean_1, sqdiff_2, mul_3 = [gf[x] for x in entry_node.outputs]
+        mean_1 = _search_nodes_by_type(gf, entry_node.outputs, 'Mean')
+        sqdiff_2 = _search_nodes_by_type(gf, entry_node.outputs, 'SquaredDifference')
+        mul_3 = _search_nodes_by_type(gf, entry_node.outputs, 'Mul')
+
         if not (mean_1.op == 'Mean' and sqdiff_2.op == 'SquaredDifference' and
             mul_3.op == 'Mul'):
             return None
@@ -284,9 +293,11 @@ def _axes_in_range(axes, rank):
             return None
         const_11 = gf[mul_10.inputs[1]]
         params['gamma'] = const_11.value.val
-        if not (gf[mul_10.outputs[0]] == mul_3 and len(mul_10.outputs) == 2):
+        if not (mul_3.name in mul_10.outputs and len(mul_10.outputs) == 2):
             return None
-        mul_12 = gf[mul_10.outputs[1]]
+        mul_12 = gf[mul_10.outputs[1]] if gf[mul_10.outputs[0]] == mul_3 else \
+            gf[mul_10.outputs[0]]
+
         sub_13 = gf[mul_12.outputs[0]]
         if not (mul_12.op == 'Mul' and sub_13.op == 'Sub'):
             return None
@@ -303,7 +314,7 @@ def _axes_in_range(axes, rank):
             add_15]
 
         return (layernorm_nodes, params)
-    except:
+    except Exception as e:
         return None
 
 
@@ -357,7 +368,10 @@ def _match_gelu_pattern(gf, entry_node):
     try:
         if not len(entry_node.outputs) == 3:
             return None
-        pow_1, add_2, mul_3 = [gf[x] for x in entry_node.outputs]
+        pow_1 = _search_nodes_by_type(gf, entry_node.outputs, 'Pow')
+        add_2 = _search_nodes_by_type(gf, entry_node.outputs, 'Add')
+        mul_3 = _search_nodes_by_type(gf, entry_node.outputs, 'Mul')
+
         if not (pow_1.op == 'Pow' and add_2.op == 'Add' and mul_3.op == 'Mul'):
             return None
         const_4 = gf[pow_1.inputs[1]]

coremltools/converters/nnssa/coreml/ssa_converter.py

@@ -1675,7 +1675,7 @@ def _convert_gelu(self, node):
             name=node.name,
             input_name=input_names[0],
             output_name=node.name,
-            mode='EXACT')
+            mode='TANH_APPROXIMATION')
 
         output_shape = self._get_tensor_shape_from_type(node.datatype)
         shapes.propagate_single_layer(layer, self.tensor_shapes,
@@ -1786,58 +1786,6 @@ def _convert_resize_nearest_neighbor(self, node):
                                       output_shapes=[output_shape])
 
 
-    def _convert_layer_normalization(self, node):
-        assert len(node.inputs) == 1
-        input_nodes, input_names, input_types = self._get_input_tensors(node)
-        input_name = input_names[0]
-        builder = self._get_builder()
-        gamma = node.attr['gamma']
-        beta = node.attr['beta']
-        axes = node.attr['axes']
-        epsilon = node.attr['epsilon']
-        input_shape = list(input_types[0].get_shape())
-
-        if (len(input_shape) in [2,3] and len(axes) == 1 and \
-            axes[0] == len(input_shape) - 1):
-            # Performance enhancement for some models with layer-norm
-            builder.add_reshape_static(name=input_name + '_reshape',
-                input_name=input_name,
-                output_name=input_name + '_reshape',
-                output_shape=input_shape + [1,1])
-
-            builder.add_mvn(name=input_name + '_mvn',
-                input_name=input_name + '_reshape',
-                output_name=input_name + '_mvn', across_channels=True,
-                normalize_variance=True, epsilon=epsilon)
-
-            builder.add_scale(name=node.name + '_5d',
-                input_name=input_name + '_mvn',
-                output_name=node.name + '_5d', W=gamma, b=beta, has_bias=True,
-                shape_scale=[len(gamma)], shape_bias=[len(beta)])
-
-            builder.add_reshape_static(name=node.name,
-                input_name=node.name + '_5d',
-                output_name=node.name,
-                output_shape=input_shape)
-
-        else:
-            # General implementation
-            input_shape = input_types[0].get_shape()
-            rdims = len(axes)
-            normalized_shape = node.datatype.get_shape()[-rdims:]
-            if gamma.shape != normalized_shape:
-                gamma = np.zeros(normalized_shape) + gamma
-            if beta.shape != normalized_shape:
-                beta = np.zeros(normalized_shape) + beta
-
-            builder.add_layer_normalization(node.name, input_name, node.name,
-                                        normalized_shape, gamma, beta, eps=1e-5)
-        
-        self.tensor_shapes[node.name] = self._get_tensor_shape_from_type(
-            node.datatype)
-
-
-
     def _convert_layer_normalization(self, node):
         assert len(node.inputs) == 1
         input_nodes, input_names, input_types = self._get_input_tensors(node)