[fix] some bugs caused by bumping

Author: Li Wei

mppq/api/interface.py

@@ -169,7 +169,7 @@ def dispatch_graph(
     quantize_operations: Collection[str] = DEFAULT_QUANTIZE_OP,
     dispatcher: Optional[str | GraphDispatcher] = None,
     dispatching_override: Optional[Dict[str, TargetPrecision]] = None,
-    ignored_scope: Optional[list | IgnoredScope] = None,
+    ignored_scope: Optional[dict | list | IgnoredScope] = None,
     quant_precision: TargetPrecision = TargetPrecision.INT8,
     **kwargs,
 ) -> BaseGraph:
@@ -216,6 +216,8 @@ def dispatch_graph(
     )
 
     if ignored_scope is not None:
+        if isinstance(ignored_scope, dict):
+            ignored_scope = IgnoredScope(**ignored_scope.pop("type"))
         if isinstance(ignored_scope, list):
             ignored_scope = IgnoredScope(operations=ignored_scope)
         assert isinstance(ignored_scope, IgnoredScope)

mppq/data.py

@@ -58,7 +58,7 @@ def from_numpy(cls, dtype: np_type):
             np_type("float16"): DataType.FP16,
             np_type("float32"): DataType.FP32,
             np_type("float64"): DataType.FP64,
-            np.bool: DataType.BOOL,
+            np.bool_: DataType.BOOL,
             np.uint8: DataType.UINT8,
             np.int8: DataType.INT8,
             np.int16: DataType.INT16,

mppq/executor/torch.py

@@ -310,10 +310,18 @@ def post_forward_hook(self, outputs: list, **kwargs) -> list:
             you can not select gpu to executing yet,
             graph will always be send to the very first visible cuda device.
         ]. Defaults to 'cuda'.
+
+        feedback_tensors=[('output_1', 'input_2')]
+        # iter 1, executor({input_2: init_tensor_2})
+        # iter 2, input_2 will be omitted and use output_1's value.
     """
 
     def __init__(
-        self, graph: BaseGraph, fp16_mode: bool = True, device: str = "cuda"
+        self,
+        graph: BaseGraph,
+        fp16_mode: bool = True,
+        device: str = "cuda",
+        feedback_tensors: dict | None = None,
     ) -> None:
         self._default_quant_fn = ppq_fake_quant
         self._deployed = False
@@ -325,6 +333,16 @@ def __init__(
         # fp16 is not available for now.
         self.fp16_mode = fp16_mode
         self.deploy()
+        if feedback_tensors:
+            self.feedback_tensors = {}
+            self.feedback_dict = {}
+            for src, dst in feedback_tensors:
+                if src not in graph.outputs:
+                    raise ValueError(f"{src} is not an output of the graph.")
+                if dst not in graph.inputs:
+                    raise ValueError(f"{dst} is not an input of the graph.")
+                self.feedback_tensors[src] = None
+                self.feedback_dict[dst] = src
 
     def register_quantize_delegate(
         self, config: TensorQuantizationConfig, delegator: TorchQuantizeDelegator
@@ -505,6 +523,11 @@ def _forward_operations(  # noqa: C901
         hooks: Optional[Mapping[str, RuntimeHook]] = None,
     ) -> List[torch.Tensor]:
         for key, value in inputs.items():
+            if value is None:
+                assert key in self.feedback_dict
+                feedback_value = self.feedback_tensors[self.feedback_dict[key]]
+                self._graph.inputs[key].value = feedback_value
+                continue
             if not isinstance(value, torch.Tensor):
                 raise TypeError(
                     "TorchExecutor can only accept tensor as its input, "
@@ -622,6 +645,12 @@ def _forward_operations(  # noqa: C901
                         result_collector[output_names.index(output_var.name)] = outputs[
                             output_idx
                         ]
+                    # collect feedback tensors
+                    if (
+                        hasattr(self, "feedback_dict")
+                        and output_var.name in self.feedback_dict.values()
+                    ):
+                        self.feedback_tensors[output_var.name] = outputs[output_idx]
             except Exception as e:
                 raise RuntimeError(f"Op Execution Error: {str(operation)}") from e
 

mppq/quantizer/base.py

@@ -104,6 +104,8 @@ def quantize(
             graph=self._graph,
             dataloader=calib_dataloader,
             executor=executor,
+            collate_fn=collate_fn,
+            calib_steps=calib_steps,
             verbose=self._verbose,
         )