Add conv clip conv pattern checking when replacing clip

TrainingExtensions/onnx/src/python/aimet_onnx/utils.py

@@ -197,39 +197,47 @@ def replace_relu6_with_relu(model: onnx_pb.ModelProto):
     :param model: ONNX model
     """
     for node in model.model.graph.node:
-        if node.op_type == 'Clip':
+        if node.op_type == 'Clip' and check_if_clip_node_minimum_is_zero(node, model):
             parent_node = None
+            child_node = None
             for temp_node in model.model.graph.node:
                 if node.input[0] in temp_node.output:
                     parent_node = temp_node
+                if node.output[0] in temp_node.input:
+                    child_node = temp_node
             assert parent_node, "Parent Node for Clip operation does not exist"
-            name = node.name
-            remove_node(node, model.model.graph)
-            inputs = [parent_node.output[0]]
-            model.replace_input_of_all_nodes(parent_node.output[0], parent_node.output[0] + '_replaced')
-            relu_node = onnx.helper.make_node(
-                op_type="Relu",
-                inputs=inputs,
-                outputs=[parent_node.output[0] + '_replaced'],
-                name='Relu_' + name,
-            )
-
-            model.add_node(relu_node)
-
-
-def check_if_node_is_relu6(node: onnx_pb.NodeProto, model: onnx_pb.ModelProto):
+            if parent_node.op_type in ['Conv', 'ConvTranspose'] and child_node and \
+                    child_node.op_type in ['Conv', 'ConvTranspose']:
+                name = node.name
+                remove_node(node, model.model.graph)
+                inputs = [parent_node.output[0]]
+                model.replace_input_of_all_nodes(parent_node.output[0], parent_node.output[0] + '_replaced')
+                relu_node = onnx.helper.make_node(
+                    op_type="Relu",
+                    inputs=inputs,
+                    outputs=[parent_node.output[0] + '_replaced'],
+                    name='Relu_' + name,
+                )
+
+                model.add_node(relu_node)
+
+
+def check_if_clip_node_minimum_is_zero(node: onnx_pb.NodeProto, model: onnx_pb.ModelProto):
     """
-    Check if the node is Relu6 by checking if it's 3rd input is a constant with value == 6
+    Check if the clip node's minimum is 0
 
     :param node: ONNX node
     :param model: ONNX model
     """
-    input_node = node.input[2]
-    for node_graph in model.model.graph.node:
-        if node_graph.output[0] == input_node:
-            if hasattr(node_graph, "attribute") and hasattr(node_graph.attribute[0], "t") and \
-                    numpy_helper.to_array(node_graph.attribute[0].t) == 6:
-                return True
+    if len(node.input) == 3:
+        input_node = node.input[1]
+        for node_graph in model.model.graph.node:
+            if node_graph.output[0] == input_node:
+                if hasattr(node_graph, "attribute") and hasattr(node_graph.attribute[0], "t") and \
+                        numpy_helper.to_array(node_graph.attribute[0].t) == 0:
+                    return True
+    elif hasattr(node, "attribute") and node.attribute[1].name == "min" and node.attribute[1].f == 0.0:
+        return True
     return False
 
 

TrainingExtensions/onnx/test/python/models/mobilenet.py

@@ -434,6 +434,7 @@ def __init__(self):
         def conv_bn(inp, oup, stride):
             return nn.Sequential(
                 nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+                nn.Hardswish(inplace=True),
                 nn.BatchNorm2d(oup),
             )
 

TrainingExtensions/onnx/test/python/test_utils.py

@@ -176,4 +176,6 @@ def test_replace_relu6_with_relu(self):
             if node.op_type == 'Relu':
                 relu_count += 1
 
-        assert relu_count - original_relu_count == relu6_count
+        hard_swish_count = 1
+
+        assert relu_count - original_relu_count == relu6_count - hard_swish_count