Fixed the best qmodel recovery issue (#1620)

yiliu30 · web-flow · commit f2d9b78b1cd5 · 2024-02-22T13:34:20.000+08:00
Signed-off-by: yiliu30 &lt;yi4.liu@intel.com&gt;
diff --git a/neural_compressor/config.py b/neural_compressor/config.py
@@ -652,7 +652,7 @@ def timeout(self):
     @timeout.setter
     def timeout(self, timeout):
         """Set timeout."""
-        if _check_value("timeout", timeout, int):
+        if _check_value("timeout", timeout, (int, float)):
             self._timeout = timeout
 
     @property
diff --git a/neural_compressor/mix_precision.py b/neural_compressor/mix_precision.py
@@ -165,7 +165,7 @@ def fit(model, conf, eval_func=None, eval_dataloader=None, eval_metric=None, **k
 
         traceback.print_exc()
     finally:
-        if strategy.best_qmodel:
+        if strategy.get_best_qmodel():
             logger.info(
                 "Specified timeout or max trials is reached! " "Found a quantized model which meet accuracy goal. Exit."
             )
@@ -176,4 +176,4 @@ def fit(model, conf, eval_func=None, eval_dataloader=None, eval_metric=None, **k
                 "Not found any quantized model which meet accuracy goal. Exit."
             )
 
-        return strategy.best_qmodel
+        return strategy.get_best_qmodel()
diff --git a/neural_compressor/quantization.py b/neural_compressor/quantization.py
@@ -249,7 +249,7 @@ def update(d):
 
         traceback.print_exc()
     finally:
-        if strategy.best_qmodel:
+        if strategy.get_best_qmodel():
             logger.info(
                 "Specified timeout or max trials is reached! " "Found a quantized model which meet accuracy goal. Exit."
             )
@@ -260,4 +260,4 @@ def update(d):
                 "Not found any quantized model which meet accuracy goal. Exit."
             )
 
-        return strategy.best_qmodel
+        return strategy.get_best_qmodel()
diff --git a/neural_compressor/strategy/auto.py b/neural_compressor/strategy/auto.py
@@ -104,7 +104,7 @@ def sequential_traverse(self):
 
             pre_strategy = strategy
             strategy.traverse()
-            self.best_qmodel = strategy.best_qmodel
+            self.best_qmodel = strategy.get_best_qmodel()
             if self.best_qmodel:
                 return
 
diff --git a/neural_compressor/strategy/strategy.py b/neural_compressor/strategy/strategy.py
@@ -1045,11 +1045,9 @@ def _eval_baseline(self):
             self.show_baseline_info()
 
     def _recover_best_qmodel_from_tuning_cfg(self):
-        """Recover the best quantized model from tuning config."""
         if self.best_tuning_cfg and not self.best_qmodel:
             logger.info(
-                f"[Strategy] Recover the {self.best_tuning_cfg.get('trial_number', 'N/A')}-trial\
-                as the tuning result."
+                "[Strategy] Recover the %s-trial as the tuning result.", self.best_tuning_cfg.get("trial_number", "N/A")
             )
             self.best_qmodel = self.adaptor.quantize(
                 copy.deepcopy(self.best_tuning_cfg), self.model, self.calib_dataloader, self.q_func
@@ -1797,6 +1795,11 @@ def __setstate__(self, d):
         """
         self.__dict__.update(d)
 
+    def get_best_qmodel(self):
+        """Get the best quantized model."""
+        self._recover_best_qmodel_from_tuning_cfg()
+        return self.best_qmodel
+
     def stop(self, timeout, trials_count):
         """Check if need to stop traverse.
 
diff --git a/neural_compressor/strategy/utils/tuning_space.py b/neural_compressor/strategy/utils/tuning_space.py
@@ -279,12 +279,6 @@ def _merge_optype_wise_cfg(self, cap: Dict, optype_wise_usr_cfg: Dict, fw_cap: D
                     cap["op"][op_name_type], op_user_cfg, fw_cap["op"][op_name_type]
                 )
 
-    def _merge_model_wise_cfg(self, cap: Dict, model_wise_usr_cfg: Dict, fw_cap: Dict):
-        for op_name_type in cap["op"].keys():
-            cap["op"][op_name_type] = self._merge_op_cfg(
-                cap["op"][op_name_type], model_wise_usr_cfg, fw_cap["op"][op_name_type]
-            )
-
     def _merge_op_wise_cfg(self, cap: Dict, op_wise_usr_cfg: Dict, fw_cap: Dict):
         op_name_types = {key[0]: key for key in cap["op"].keys()}
         for op_name_pattern, op_user_cfg in op_wise_usr_cfg.items():
@@ -406,8 +400,6 @@ def _merge_with_user_cfg(self, capability: Dict, user_cfg: Dict):
         :return:
         """
         fw_capability = deepcopy(capability)
-        if user_cfg["model_wise"] is not None:
-            self._merge_model_wise_cfg(capability, user_cfg["model_wise"], fw_capability)
         if user_cfg["optype_wise"] is not None:
             self._merge_optype_wise_cfg(capability, user_cfg["optype_wise"], fw_capability)
         if user_cfg["op_wise"] is not None:
diff --git a/neural_compressor/training.py b/neural_compressor/training.py
@@ -341,7 +341,7 @@ def eval_func(model):
 
         traceback.print_exc()
     finally:
-        if strategy.best_qmodel:
+        if strategy.get_best_qmodel():
             logger.info(
                 "Specified timeout or max trials is reached! " "Found a quantized model which meet accuracy goal. Exit."
             )
@@ -352,7 +352,7 @@ def eval_func(model):
                 "Not found any quantized model which meet accuracy goal. Exit."
             )
 
-        compression_manager.model = strategy.best_qmodel
+        compression_manager.model = strategy.get_best_qmodel()
 
     return compression_manager.model
 

Original file line number	Diff line number	Diff line change
`@@ -165,7 +165,7 @@ def fit(model, conf, eval_func=None, eval_dataloader=None, eval_metric=None, **k`
`165`	`165`
`166`	`166`	`traceback.print_exc()`
`167`	`167`	`finally:`
`168`		`- if strategy.best_qmodel:`
	`168`	`+ if strategy.get_best_qmodel():`
`169`	`169`	`logger.info(`
`170`	`170`	`"Specified timeout or max trials is reached! " "Found a quantized model which meet accuracy goal. Exit."`
`171`	`171`	`)`
`@@ -176,4 +176,4 @@ def fit(model, conf, eval_func=None, eval_dataloader=None, eval_metric=None, **k`
`176`	`176`	`"Not found any quantized model which meet accuracy goal. Exit."`
`177`	`177`	`)`
`178`	`178`
`179`		`- return strategy.best_qmodel`
	`179`	`+ return strategy.get_best_qmodel()`
Original file line number	Diff line number	Diff line change
`@@ -249,7 +249,7 @@ def update(d):`
`249`	`249`
`250`	`250`	`traceback.print_exc()`
`251`	`251`	`finally:`
`252`		`- if strategy.best_qmodel:`
	`252`	`+ if strategy.get_best_qmodel():`
`253`	`253`	`logger.info(`
`254`	`254`	`"Specified timeout or max trials is reached! " "Found a quantized model which meet accuracy goal. Exit."`
`255`	`255`	`)`
`@@ -260,4 +260,4 @@ def update(d):`
`260`	`260`	`"Not found any quantized model which meet accuracy goal. Exit."`
`261`	`261`	`)`
`262`	`262`
`263`		`- return strategy.best_qmodel`
	`263`	`+ return strategy.get_best_qmodel()`
Original file line number	Diff line number	Diff line change
`@@ -341,7 +341,7 @@ def eval_func(model):`
`341`	`341`
`342`	`342`	`traceback.print_exc()`
`343`	`343`	`finally:`
`344`		`- if strategy.best_qmodel:`
	`344`	`+ if strategy.get_best_qmodel():`
`345`	`345`	`logger.info(`
`346`	`346`	`"Specified timeout or max trials is reached! " "Found a quantized model which meet accuracy goal. Exit."`
`347`	`347`	`)`
`@@ -352,7 +352,7 @@ def eval_func(model):`
`352`	`352`	`"Not found any quantized model which meet accuracy goal. Exit."`
`353`	`353`	`)`
`354`	`354`
`355`		`- compression_manager.model = strategy.best_qmodel`
	`355`	`+ compression_manager.model = strategy.get_best_qmodel()`
`356`	`356`
`357`	`357`	`return compression_manager.model`
`358`	`358`