Canonicalize (un)squeeze operations as pre-opset-13, i.e., put axes into node attributes to ommit creating a buffer for it

Author: lukamac

Deeploy/OperatorDescriptor.py

@@ -476,15 +476,29 @@ def canonicalize(self, node: gs.Node, opset: int) -> bool:
     attrDescriptors = [AttrDesc("axis", IntUnpack, default = 0)],
 )
 
+
+class SqueezeDescriptor(OperatorDescriptor):
+
+    def canonicalize(self, node: gs.Node, opset: int) -> bool:
+        if opset >= 13:
+            assert len(node.inputs) == 2, f"Expected 2 inputs but received {len(node.inputs)}"
+            axes = node.inputs[1]
+            assert isinstance(axes,
+                              gs.Constant), f"Expected axes to be a constant but received axes of type {type(axes)}"
+            node.attrs["axes"] = axes.values
+            axes.outputs.clear()
+        return super().canonicalize(node, opset)
+
+
 # Opset <= 11
-unsqueezeDesc = OperatorDescriptor(
+unsqueezeDesc = SqueezeDescriptor(
     inputDescriptor = IoDesc("data_in"),
     outputDescriptor = IoDesc("data_out"),
     attrDescriptors = [AttrDesc("axes", IntTupleUnpack)],
 )
 
 # Opset <= 11
-squeezeDesc = OperatorDescriptor(
+squeezeDesc = SqueezeDescriptor(
     inputDescriptor = IoDesc("data_in"),
     outputDescriptor = IoDesc("data_out"),
     attrDescriptors = [AttrDesc("axes", IntTupleUnpack)],

Deeploy/Targets/Generic/Parsers.py

@@ -986,48 +986,23 @@ def __init__(self):
         super().__init__()
 
     def parseNode(self, node: gs.Node) -> (bool):
+        if not all(['axes' in node.attrs, len(node.inputs) == 1, len(node.outputs) == 1]):
+            return False
 
-        # ONNX v11: 'axes' is a node attribute
-        if 'axes' in node.attrs:
-            ret = all(['axes' in node.attrs, len(node.inputs) == 1, len(node.outputs) == 1])
-        # ONNX v13+: 'axes' becomes an input with the data
-        # Source: https://onnx.ai/onnx/operators/onnx__Unsqueeze.html
-        else:
-            ret = all([len(node.inputs) == 2, len(node.outputs) == 1])
-
-        if ret and 'axes' in node.attrs:
-            axes_attr = node.attrs['axes']
-            self.operatorRepresentation['axes'] = [int(axes_attr)] if isinstance(axes_attr, int) \
-                else [int(a) for a in axes_attr]
-        # For opset 13+, axes will be extracted from the second input in parseNodeCtxt
-
-        return ret
+        self.operatorRepresentation['axes'] = node.attrs['axes']
+        return True
 
     def parseNodeCtxt(self,
                       ctxt: NetworkContext,
                       node: gs.Node,
                       channels_first: bool = True) -> Tuple[NetworkContext, bool]:
+        inputs = ['data_in']
+        for idx, inputNode in enumerate(node.inputs):
+            self.operatorRepresentation[inputs[idx]] = ctxt.lookup(inputNode.name).name
 
         outputs = ['data_out']
-        if len(node.inputs) == 1:
-            inputs = ['data_in']
-            for idx, inputNode in enumerate(node.inputs):
-                self.operatorRepresentation[inputs[idx]] = ctxt.lookup(inputNode.name).name
-            for idx, outputNode in enumerate(node.outputs):
-                self.operatorRepresentation[outputs[idx]] = ctxt.lookup(outputNode.name).name
-        else:
-            data_in = ctxt.lookup(node.inputs[0].name)
-            data_out = ctxt.lookup(node.outputs[0].name)
-            self.operatorRepresentation['data_in'] = data_in.name
-            self.operatorRepresentation['data_out'] = data_out.name
-            # axes must be a constant; extract values
-            axes_buf = ctxt.lookup(node.inputs[1].name)
-            assert hasattr(axes_buf, 'values'), "Unsqueeze: expected constant 'axes' input for opset 13+"
-            axes_vals = np.array(axes_buf.values).astype(int).flatten().tolist()
-            self.operatorRepresentation['axes'] = axes_vals
-            # Do not deploy the axes tensor
-            axes_buf._live = False
-            axes_buf._deploy = False
+        for idx, outputNode in enumerate(node.outputs):
+            self.operatorRepresentation[outputs[idx]] = ctxt.lookup(outputNode.name).name
 
         return ctxt, True