[fix] fix several glitches

- torch executor forward with None is accepted
- add back trainable graph
- set master_by correctly

Author: LoSealL

mppq/__init__.py

@@ -7,4 +7,4 @@
 Minimized PPQ quantizer package.
 """
 
-__version__ = "0.7.2"
+__version__ = "0.7.3"

mppq/api/__init__.py

@@ -14,6 +14,7 @@
 from mppq.api.extension import register_operation, register_platform
 from mppq.api.interface import (
     dispatch_graph,
+    export_config,
     export_graph,
     export_onnx_graph,
     format_graph,
@@ -37,6 +38,7 @@
     "register_operation",
     "register_platform",
     "dispatch_graph",
+    "export_config",
     "export_graph",
     "export_onnx_graph",
     "format_graph",

mppq/api/extension.py

@@ -86,10 +86,12 @@ def register_platform(
             字典键值可选，作为调度器的命名，和调度器的类类型。
         quantizer (Dict[str | None, Type[BaseQuantizer]]): 自定义平台的量化器。
             字典键值可选，作为量化器的命名，和量化器的类类型。
-        parsers (Optional[Dict[str | None, Type[GraphBuilder]]], optional): 自定义平台的图构建器。
-            字典键值可选，作为图构建器的命名，和图构建器的类类型。 Defaults to None.
-        exporters (Optional[Dict[str | None, Type[GraphExporter]]], optional): 自定义平台的图导出器。
-            字典键值可选，作为图导出器的命名，和图导出器的类类型。 Defaults to None.
+        parsers (Optional[Dict[str | None, Type[GraphBuilder]]], optional):
+            自定义平台的图构建器。字典键值可选，作为图构建器的命名，和图构建器的类类型。
+            Defaults to None.
+        exporters (Optional[Dict[str | None, Type[GraphExporter]]], optional):
+            自定义平台的图导出器。字典键值可选，作为图导出器的命名，和图导出器的类类型。
+            Defaults to None.
     """
     if platform_id in _PLATFORM_TO_DISPATCHER_:
         raise KeyError(f"Platform {platform_id} is already registered.")

mppq/api/interface.py

@@ -86,6 +86,17 @@ def export_onnx_graph(graph: BaseGraph, f: str | os.PathLike):
     export_graph(graph, f, exporter=exporter)
 
 
+def export_config(graph: BaseGraph, f: str | os.PathLike):
+    r"""导出PPQ IR Graph的量化配置。
+    Args:
+        graph (BaseGraph): 待导出IR Graph对象。
+        f (str|os.PathLike): 导出文件路径。
+    """
+
+    exporter = EXPORTER["onnx"]()
+    exporter.dump_quantization_config(f, graph)
+
+
 def format_graph(
     graph: BaseGraph,
     format_constant_input: bool = True,

mppq/data.py

@@ -3,7 +3,7 @@
 You are not allowed to modify this 请勿修改此文件
 """
 
-from enum import Enum
+from enum import IntEnum
 from numbers import Number
 from typing import Any, Literal, Optional, Sequence, overload
 
@@ -14,7 +14,7 @@
 from torch import dtype as torch_type
 
 
-class DataType(Enum):
+class DataType(IntEnum):
     """
     DataType defines all PPQ internal data type and its enumeration value.
         ATTENTION: PPQ shares same data type enumeration value with Onnx.
@@ -113,8 +113,8 @@ def to_numpy(cls, dtype) -> np_type:
             DataType.FP32: np_type("float32"),
             DataType.FP64: np_type("float64"),
         }
-        assert isinstance(dtype, DataType)
-        return numpy_converting_dict[dtype]
+        assert isinstance(dtype, int)
+        return numpy_converting_dict[DataType(dtype)]
 
     @classmethod
     def to_torch(cls, dtype) -> torch_type:
@@ -131,8 +131,8 @@ def to_torch(cls, dtype) -> torch_type:
             DataType.FP64: torch.float64,
             DataType.FP8_E5M2: torch.float8_e5m2,
         }
-        assert isinstance(dtype, DataType)
-        return torch_converting_dict[dtype]
+        assert isinstance(dtype, int)
+        return torch_converting_dict[DataType(dtype)]
 
 
 @overload

mppq/executor/base.py

@@ -46,7 +46,9 @@ class RuntimeHook:
     def __init__(self, operation: Operation, **kwargs) -> None:
         self._hook_to = operation
 
-    def pre_forward_hook(self, inputs: Sequence[Tensor], **kwargs) -> List[Tensor]:
+    def pre_forward_hook(
+        self, inputs: Sequence[Tensor | None], **kwargs
+    ) -> List[Tensor | None]:
         """user-customized pre-processing procedure of input data.
 
         Args:
@@ -57,7 +59,9 @@ def pre_forward_hook(self, inputs: Sequence[Tensor], **kwargs) -> List[Tensor]:
         """
         return list(inputs)
 
-    def post_forward_hook(self, outputs: Sequence[Tensor], **kwargs) -> List[Tensor]:
+    def post_forward_hook(
+        self, outputs: Sequence[Tensor | None], **kwargs
+    ) -> List[Tensor | None]:
         """user-customized post-processing procedure of output data.
 
         Args:
@@ -82,21 +86,21 @@ def __init__(self, operation: Operation, **kwargs) -> None:
 
     def pre_forward_hook(
         self,
-        inputs: Sequence[Tensor],
-        quant_inputs: Sequence[Tensor] = (),
+        inputs: Sequence[Tensor | None],
+        quant_inputs: Sequence[Tensor | None] = (),
         quant_configs: Sequence[TensorQuantizationConfig] = (),
         **kwargs,
-    ) -> List[Tensor]:
+    ) -> List[Tensor | None]:
         assert len(inputs) == len(quant_inputs) == len(quant_configs)
         return list(quant_inputs)
 
     def post_forward_hook(
         self,
-        outputs: Sequence[Tensor],
-        quant_outputs: Sequence[Tensor] = (),
+        outputs: Sequence[Tensor | None],
+        quant_outputs: Sequence[Tensor | None] = (),
         quant_configs: Sequence[TensorQuantizationConfig] = (),
         **kwargs,
-    ) -> List[Tensor]:
+    ) -> List[Tensor | None]:
         assert len(outputs) == len(quant_outputs) == len(quant_configs)
         return list(quant_outputs)
 
@@ -155,7 +159,7 @@ def _prepare_input(self, inputs: Optional[GraphInput]) -> Dict[str, Tensor]:
     def forward(
         self,
         inputs: GraphInput,
-        output_names: Optional[List[str]] = None,
+        output_names: Optional[Sequence[str]] = None,
         hooks: Optional[Mapping[str, RuntimeHook]] = None,
     ) -> List[torch.Tensor]:
         """Forward a graph from given inputs to required output names.
@@ -169,15 +173,15 @@ def forward(
     def tracing_operation_meta(
         self,
         inputs: GraphInput,
-        output_names: Optional[List[str]] = None,
+        output_names: Optional[Sequence[str]] = None,
     ) -> None:
         raise NotImplementedError("Please implement this function first.")
 
     @overload
     def forward_single_operation(
         self,
         op: Operation,
-        inputs: List[Tensor],
+        inputs: Sequence[Tensor | None],
         ctx: Optional[TorchBackendContext] = None,
         return_list: Literal[True] = True,
     ) -> Tuple[Tensor, ...]:
@@ -187,7 +191,7 @@ def forward_single_operation(
     def forward_single_operation(
         self,
         op: Operation,
-        inputs: List[Tensor],
+        inputs: Sequence[Tensor | None],
         ctx: Optional[TorchBackendContext] = None,
         return_list: bool = True,
     ) -> Tuple[Tensor, ...] | Tensor:
@@ -196,7 +200,7 @@ def forward_single_operation(
     def forward_single_operation(
         self,
         op: Operation,
-        inputs: List[Tensor],
+        inputs: Sequence[Tensor | None],
         ctx: Optional[TorchBackendContext] = None,
         return_list: bool = True,
     ) -> Tuple[Tensor, ...] | Tensor:
@@ -210,7 +214,7 @@ def forward_single_operation(
     def __call__(
         self,
         inputs: GraphInput,
-        output_names: Optional[List[str]] = None,
+        output_names: Optional[Sequence[str]] = None,
     ) -> List[torch.Tensor]:
         return self.forward(inputs=inputs, output_names=output_names)
 

mppq/executor/op/base.py

@@ -116,7 +116,7 @@ class OperationForwardProtocol(Protocol):
     def __call__(
         self,
         op: Operation,
-        values: Sequence[torch.Tensor],
+        values: Sequence[torch.Tensor | None],
         ctx: Optional[TorchBackendContext] = None,
         **kwargs,
     ) -> torch.Tensor | Tuple[torch.Tensor, ...]: