Onnx slim transform

QEfficient/base/modeling_qeff.py

@@ -45,9 +45,10 @@
     def _transform_names(cls) -> List[str]:
         return [x.__name__ for x in cls._pytorch_transforms + cls._onnx_transforms]
 
-    def __init__(self, model: torch.nn.Module) -> None:
+    def __init__(self, model: torch.nn.Module, onnx_slim_transfom: bool = False) -> None:
         super().__init__()
         self.model = model
+        self.onnx_slim_transform = onnx_slim_transfom
         self.onnx_path: Optional[str] = None
         self.qpc_path: Optional[str] = None
         self.qpc_session: Optional[QAICInferenceSession] = None
@@ -119,6 +120,7 @@
         example_inputs: Dict[str, torch.Tensor],
         output_names: List[str],
         dynamic_axes: Dict[str, Dict[int, str]],
+        onnx_slim_transform: bool = False,
         export_kwargs: Optional[Dict[str, any]] = None,
         onnx_transform_kwargs: Optional[Dict[str, any]] = None,
         export_dir: Optional[str] = None,
@@ -146,7 +148,6 @@
         tmp_onnx_dir.mkdir(parents=True, exist_ok=True)
 
         # Create input_names from example_inputs
-
         input_names = []
         for param in inspect.signature(self.model.forward).parameters:
             if param in example_inputs:
@@ -183,11 +184,14 @@
                 **export_kwargs,
             )
             logger.info("Pytorch export successful")
-
             model = onnx.load(tmp_onnx_path, load_external_data=False)
             transform_kwargs = {
+                "temp_onnx_path": tmp_onnx_path,
+                "model_name": self.model_name,
+                "enable_onnx_slim_transform": onnx_slim_transform,
                 "onnx_base_dir": str(tmp_onnx_dir),
                 "model_name": self.model_name,
+                "enable_onnx_slim_transform": onnx_slim_transform,
             }
             if onnx_transform_kwargs is not None:
                 transform_kwargs.update(onnx_transform_kwargs)
@@ -249,7 +253,6 @@
         """
         if onnx_path is None and self.onnx_path is None:
             self.export()
-
         onnx_path = Path(onnx_path or self.onnx_path)
         compile_dir = Path(compile_dir or onnx_path.parent)
         qpc_path = compile_dir / "qpc"
@@ -368,5 +371,4 @@
             )
 
         self.qpc_path = qpc_path
-
         return qpc_path

QEfficient/base/onnx_transforms.py

@@ -8,6 +8,8 @@
 from typing import Optional, Tuple
 
 import numpy as np
+import onnx
+import onnxslim
 from onnx import ModelProto, external_data_helper, numpy_helper
 
 
@@ -37,6 +39,8 @@ class FP16ClipTransform(OnnxTransform):
     Clips the tensor values to be in FP16 range, but preserves -inf values.
     """
 
+    print("FP16ClipTransform is applied")
+
     @classmethod
     def apply(cls, model: ModelProto, *, onnx_base_dir: Optional[str] = None, **kwargs) -> Tuple[ModelProto, bool]:
         """
@@ -99,3 +103,32 @@ def apply(
                     current_file_size = tsize
                 external_data_helper.set_external_data(tensor, f"{model_name}_{file_num}.onnx.data")
         return model, transformed
+
+
+class OnnxSlimTransform(OnnxTransform):
+    """
+    Applies onnx-slim transformations on the given ONNX graph.
+    """
+
+    @classmethod
+    def apply(
+        cls,
+        model: ModelProto,
+        *,
+        onnx_base_dir: Optional[str] = None,
+        **kwargs,
+    ) -> Tuple[ModelProto, bool]:
+        """
+        :param enable_onnx_slim_transform: If True, applies onnx-slim transformations.
+        """
+        # print(kwargs)
+        transformed = False
+        onnx_slim_transform = kwargs.get("enable_onnx_slim_transform", False)
+        temp_onnx_path = kwargs.get("temp_onnx_path", None)
+        if onnx_slim_transform:
+            print("onnx slim transform done")
+            transformed = True
+            slimmed_model = onnxslim.slim(model)
+            onnx.save(slimmed_model, temp_onnx_path)
+            return slimmed_model, transformed
+        return model, transformed