Merge pull request cvg#46 from fabio-sim/fabio/feat/cpu-optim

feat: CPU compatibility (cvg#46)

Author: fabio-sim

lightglue_onnx/optim/fusion_attention_lightglue.py

@@ -211,6 +211,7 @@ def create_self_attention_node(
     ) -> NodeProto:
         # all_inputs are (B, N, S, H)
         if self.enable_packed_qkv:
+            # Implement Stack via Unsqueeze+Concat
             unsqueeze_q_node_name = self.model.create_node_name("Unsqueeze")
             unsqueeze_k_node_name = self.model.create_node_name("Unsqueeze")
             unsqueeze_v_node_name = self.model.create_node_name("Unsqueeze")
@@ -297,8 +298,98 @@ def create_self_attention_node(
             )
 
             return attention_node
-        else:  # Not packed
-            raise NotImplementedError("Unpacked QKV self-attention not implemented.")
+        else:  # Not packed. (CPU-compatible)
+            # Transpose nodes: (B, N, S, H) -> (B, S, N, H)
+            transpose_q_node_name = self.model.create_node_name("Transpose")
+            transpose_k_node_name = self.model.create_node_name("Transpose")
+            transpose_v_node_name = self.model.create_node_name("Transpose")
+            transpose_q_node = helper.make_node(
+                "Transpose",
+                inputs=[matmul_q.output[0]],
+                outputs=[transpose_q_node_name + "_out"],
+                name=transpose_q_node_name,
+                perm=[0, 2, 1, 3],
+            )
+            self.node_name_to_graph_name[transpose_q_node.name] = self.this_graph_name
+            transpose_k_node = helper.make_node(
+                "Transpose",
+                inputs=[matmul_k.output[0]],
+                outputs=[transpose_k_node_name + "_out"],
+                name=transpose_k_node_name,
+                perm=[0, 2, 1, 3],
+            )
+            self.node_name_to_graph_name[transpose_k_node.name] = self.this_graph_name
+            transpose_v_node = helper.make_node(
+                "Transpose",
+                inputs=[matmul_v.output[0]],
+                outputs=[transpose_v_node_name + "_out"],
+                name=transpose_v_node_name,
+                perm=[0, 2, 1, 3],
+            )
+            self.node_name_to_graph_name[transpose_v_node.name] = self.this_graph_name
+
+            # Reshape nodes: (B, S, N, H) -> (B, S, NH)
+            reshape_q_node_name = self.model.create_node_name("Reshape")
+            reshape_k_node_name = self.model.create_node_name("Reshape")
+            reshape_v_node_name = self.model.create_node_name("Reshape")
+            for n in (reshape_q_node_name, reshape_k_node_name, reshape_v_node_name):
+                self.add_initializer(
+                    name=n + "_shape",
+                    data_type=TensorProto.INT64,
+                    dims=[3],
+                    vals=[0, 0, hidden_size],
+                    raw=False,
+                )
+            reshape_q_node = helper.make_node(
+                "Reshape",
+                inputs=[transpose_q_node_name + "_out", reshape_q_node_name + "_shape"],
+                outputs=[reshape_q_node_name + "_out"],
+                name=reshape_q_node_name,
+            )
+            self.node_name_to_graph_name[reshape_q_node.name] = self.this_graph_name
+            reshape_k_node = helper.make_node(
+                "Reshape",
+                inputs=[transpose_k_node_name + "_out", reshape_k_node_name + "_shape"],
+                outputs=[reshape_k_node_name + "_out"],
+                name=reshape_k_node_name,
+            )
+            self.node_name_to_graph_name[reshape_k_node.name] = self.this_graph_name
+            reshape_v_node = helper.make_node(
+                "Reshape",
+                inputs=[transpose_v_node_name + "_out", reshape_v_node_name + "_shape"],
+                outputs=[reshape_v_node_name + "_out"],
+                name=reshape_v_node_name,
+            )
+            self.node_name_to_graph_name[reshape_v_node.name] = self.this_graph_name
+
+            self.nodes_to_add.extend(
+                [
+                    transpose_q_node,
+                    transpose_k_node,
+                    transpose_v_node,
+                    reshape_q_node,
+                    reshape_k_node,
+                    reshape_v_node,
+                ]
+            )
+
+            attention_inputs = [
+                reshape_q_node_name + "_out",
+                reshape_k_node_name + "_out",
+                reshape_v_node_name + "_out",
+            ]
+
+            attention_node_name = self.model.create_node_name("MultiHeadAttention")
+            attention_node = helper.make_node(
+                "MultiHeadAttention",
+                inputs=attention_inputs,
+                outputs=[output],
+                name=attention_node_name,
+                domain="com.microsoft",
+                num_heads=num_heads,
+            )
+
+            return attention_node
 
     def create_cross_attention_node(
         self,

optimize.py

@@ -5,6 +5,7 @@
 
 from onnx import load_model, save_model
 from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
+from onnxruntime.transformers.fusion_options import FusionOptions
 
 from lightglue_onnx.optim.onnx_model_lightglue import LightGlueOnnxModel
 
@@ -19,6 +20,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "-o", "--output", type=str, help="Path to output fused LightGlue ONNX model."
     )
+    parser.add_argument(
+        "--cpu", action="store_true", help="Whether to optimize for CPU."
+    )
 
     return parser.parse_args()
 
@@ -28,12 +32,19 @@ def parse_args() -> argparse.Namespace:
     lightglue = load_model(args.input)
     optimizer = LightGlueOnnxModel(lightglue, NUM_HEADS, HIDDEN_SIZE)
 
-    optimizer.optimize()
+    options = None
+    if args.cpu:
+        options = FusionOptions("unet")
+        options.enable_packed_qkv = False
+
+    optimizer.optimize(options)
     optimizer.get_fused_operator_statistics()
 
     output_path = args.output
     if output_path is None:
         output_path = args.input.replace(".onnx", "_fused.onnx")
+        if args.cpu:
+            output_path = output_path.replace(".onnx", "_cpu.onnx")
 
     optimizer.save_model_to_file(output_path)
 
@@ -42,6 +53,10 @@ def parse_args() -> argparse.Namespace:
         output_path,
     )
 
+    if args.cpu:
+        print("CPU does not support fp16. Skipping..")
+        exit()
+
     optimizer.convert_float_to_float16(
         keep_io_types=True,
     )