Change Neureka conv parsers to inherit from Conv2d parser

Author: lukamac

Deeploy/Targets/Neureka/Parsers.py

@@ -28,20 +28,16 @@
 import onnx_graphsurgeon as gs
 
 from Deeploy.DeeployTypes import NetworkContext
-from Deeploy.Targets.Generic.Parsers import ConvParser, RQSParserInterface
+from Deeploy.Targets.Generic.Parsers import Conv2DParser, ConvParser, RQSParserInterface
 
 
-class NeurekaConv2DBaseParser(ConvParser):
-
-    def __init__(self, noBiasHoisting: bool = True):
-        super().__init__(noBiasHoisting)
+class NeurekaConv2DBaseParser(Conv2DParser):
 
     def parseNode(self, node: gs.Node) -> bool:
         if not super().parseNode(node):
             return False
 
         if not all([
-                len(node.attrs['pads']) == 4,
                 # No dilation support
                 self.operatorRepresentation['dilations'] == [1, 1],
                 # Channels have to be last
@@ -51,16 +47,6 @@ def parseNode(self, node: gs.Node) -> bool:
         ]):
             return False
 
-        self.operatorRepresentation['dim_kernel_x'] = int(self.operatorRepresentation['kernel_shape'][0])
-        self.operatorRepresentation['dim_kernel_y'] = int(self.operatorRepresentation['kernel_shape'][1])
-        self.operatorRepresentation['dilation_x'] = int(self.operatorRepresentation['dilations'][0])
-        self.operatorRepresentation['dilation_y'] = int(self.operatorRepresentation['dilations'][1])
-        self.operatorRepresentation['padding_x'] = int(self.operatorRepresentation['pads'][0])
-        self.operatorRepresentation['padding_y'] = int(self.operatorRepresentation['pads'][1])
-        self.operatorRepresentation['stride_x'] = int(self.operatorRepresentation['strides'][0])
-        self.operatorRepresentation['stride_y'] = int(self.operatorRepresentation['strides'][1])
-        self.operatorRepresentation['bias_shift'] = int(0)
-        self.operatorRepresentation['out_shift'] = int(0)
         self.operatorRepresentation['padding_y_top'] = int(self.operatorRepresentation['pads'][0])
         self.operatorRepresentation['padding_x_left'] = int(self.operatorRepresentation['pads'][1])
         self.operatorRepresentation['padding_y_bottom'] = int(self.operatorRepresentation['pads'][2])
@@ -73,31 +59,36 @@ def parseNodeCtxt(self,
                       ctxt: NetworkContext,
                       node: gs.Node,
                       channels_first: bool = True) -> Tuple[NetworkContext, bool]:
-
-        newCtxt, ret = super().parseNodeCtxt(ctxt, node, channels_first)
+        # LMACAN: Cannot reuse the Conv2DParser's parserNodeCtxt because it requires the weight shape
+        #         to be of length 4 whereas neureka does a specific weight encoding so the shape
+        #         ends up being equal to 3.
+        newCtxt, ret = ConvParser.parseNodeCtxt(self, ctxt, node, channels_first)
 
         if not ret:
             return ctxt, False
 
+        # LMACAN: c/p of Conv2DParser's parserNodeCtxt but with a different weight shape check
+        #         and enforcing that the channels_first is false
         data_in = newCtxt.lookup(self.operatorRepresentation['data_in'])
         data_out = newCtxt.lookup(self.operatorRepresentation['data_out'])
         weight = newCtxt.lookup(self.operatorRepresentation['weight'])
 
+        if not all([
+                channels_first == False,
+                len(data_in.shape) == 4,
+                # LMACAN: weight shape should be equal to 3 because we have to do the neureka's
+                #         special weight encoding
+                len(weight.shape) == 3,
+        ]):
+            return newCtxt, False
+
         self.operatorRepresentation['batch'] = data_in.shape[0]
-        if channels_first:
-            self.operatorRepresentation['ch_im_in'] = data_in.shape[1]
-            self.operatorRepresentation['dim_im_in_x'] = data_in.shape[2]
-            self.operatorRepresentation['dim_im_in_y'] = data_in.shape[3]
-            self.operatorRepresentation['ch_im_out'] = data_out.shape[1]
-            self.operatorRepresentation['dim_im_out_x'] = data_out.shape[2]
-            self.operatorRepresentation['dim_im_out_y'] = data_out.shape[3]
-        else:
-            self.operatorRepresentation['ch_im_in'] = data_in.shape[3]
-            self.operatorRepresentation['dim_im_in_x'] = data_in.shape[1]
-            self.operatorRepresentation['dim_im_in_y'] = data_in.shape[2]
-            self.operatorRepresentation['ch_im_out'] = data_out.shape[3]
-            self.operatorRepresentation['dim_im_out_x'] = data_out.shape[1]
-            self.operatorRepresentation['dim_im_out_y'] = data_out.shape[2]
+        self.operatorRepresentation['dim_im_in_x'] = data_in.shape[1]
+        self.operatorRepresentation['dim_im_in_y'] = data_in.shape[2]
+        self.operatorRepresentation['ch_im_in'] = data_in.shape[3]
+        self.operatorRepresentation['dim_im_out_x'] = data_out.shape[1]
+        self.operatorRepresentation['dim_im_out_y'] = data_out.shape[2]
+        self.operatorRepresentation['ch_im_out'] = data_out.shape[3]
 
         # No requantization
         self.operatorRepresentation['mul'] = 'NULL'
@@ -113,28 +104,18 @@ def parseNode(self, node: gs.Node) -> bool:
         if not super().parseNode(node):
             return False
 
-        if not self.operatorRepresentation['kernel_shape'] == [3, 3]:
-            return False
-        if self.operatorRepresentation['group'] == 1:
+        ch_im_out = node.inputs[1].shape[0]
+        ch_im_in = node.inputs[1].shape[1]
+
+        if not all([
+                self.operatorRepresentation['kernel_shape'] == [3, 3],
+                self.operatorRepresentation['group'] == ch_im_out,
+                self.operatorRepresentation['group'] == ch_im_in,
+        ]):
             return False
 
         return True
 
-    def parseNodeCtxt(self,
-                      ctxt: NetworkContext,
-                      node: gs.Node,
-                      channels_first: bool = True) -> Tuple[NetworkContext, bool]:
-
-        newCtxt, ret = super().parseNodeCtxt(ctxt, node, channels_first)
-
-        data_in = ctxt.lookup(self.operatorRepresentation['data_in'])
-        weight = ctxt.lookup(self.operatorRepresentation['weight'])
-
-        if len(data_in.shape) != 4 or len(weight.shape) != 4:
-            return ctxt, False
-
-        return newCtxt, True
-
 
 class NeurekaRQSDWConv2DParser(NeurekaDWConv2DParser, RQSParserInterface):
 
@@ -168,29 +149,14 @@ def parseNode(self, node: gs.Node) -> bool:
         if not super().parseNode(node):
             return False
 
-        if not self.operatorRepresentation['kernel_shape'] == [1, 1]:
+        if not all([
+                self.operatorRepresentation['kernel_shape'] == [1, 1],
+                self.operatorRepresentation['group'] == 1,
+        ]):
             return False
 
-        # if not self.operatorRepresentation['strides'] == [1, 1]:
-        #     return False
-
         return True
 
-    def parseNodeCtxt(self,
-                      ctxt: NetworkContext,
-                      node: gs.Node,
-                      channels_first: bool = True) -> Tuple[NetworkContext, bool]:
-
-        newCtxt, ret = super().parseNodeCtxt(ctxt, node, channels_first)
-
-        data_in = newCtxt.lookup(self.operatorRepresentation['data_in'])
-        weight = newCtxt.lookup(self.operatorRepresentation['weight'])
-
-        if len(data_in.shape) != 4 or len(weight.shape) != 3:
-            return ctxt, False
-
-        return newCtxt, True
-
 
 class NeurekaRQSPWConv2DParser(NeurekaPWConv2DParser, RQSParserInterface):
 
@@ -223,29 +189,14 @@ def parseNode(self, node: gs.Node) -> bool:
         if not super().parseNode(node):
             return False
 
-        if not self.operatorRepresentation['kernel_shape'] == [3, 3]:
-            return False
-
-        if not self.operatorRepresentation['group'] == 1:
+        if not all([
+                self.operatorRepresentation['kernel_shape'] == [3, 3],
+                self.operatorRepresentation['group'] == 1,
+        ]):
             return False
 
         return True
 
-    def parseNodeCtxt(self,
-                      ctxt: NetworkContext,
-                      node: gs.Node,
-                      channels_first: bool = True) -> Tuple[NetworkContext, bool]:
-
-        newCtxt, ret = super().parseNodeCtxt(ctxt, node, channels_first)
-
-        data_in = newCtxt.lookup(self.operatorRepresentation['data_in'])
-        weight = newCtxt.lookup(self.operatorRepresentation['weight'])
-
-        if len(data_in.shape) != 4 or len(weight.shape) != 4:
-            return ctxt, False
-
-        return newCtxt, True
-
 
 class NeurekaRQSDenseConv2DParser(NeurekaDenseConv2DParser, RQSParserInterface):