Fixed the best qmodel recovery issue (#1620)

Signed-off-by: yiliu30 <yi4.liu@intel.com>

Author: yiliu30

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

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()

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()

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
 

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.
 

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:

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