use safe pickle_load for resume_file (#2329)

Signed-off-by: chensuyue <suyue.chen@intel.com>

Author: chensuyue

neural_compressor/adaptor/tensorflow.py

@@ -1354,7 +1354,7 @@ def inspect_tensor(
             model = model.graph_def
         if not quantization_cfg:
             # TODO get config from graph if config is None
-            quantization_cfg = load_data_from_pkl("./nc_workspace/", "cfg.pkl")
+            quantization_cfg = load_data_from_pkl("./nc_workspace/cfg.pkl")
         node_list = op_list
         # create the mapping between node name and node, key: node_name, val: node
         graph_node_name_mapping = {}

neural_compressor/mix_precision.py

@@ -29,7 +29,7 @@
 from .model import Model
 from .strategy import STRATEGIES
 from .utils import alias_param, logger
-from .utils.utility import CpuInfo, secure_check_eval_func, time_limit
+from .utils.utility import CpuInfo, load_data_from_pkl, secure_check_eval_func, time_limit
 
 
 @alias_param("conf", param_alias="config")
@@ -142,7 +142,7 @@ def fit(model, conf, eval_func=None, eval_dataloader=None, eval_metric=None, **k
     if resume_file:
         assert os.path.exists(resume_file), "The specified resume file {} doesn't exist!".format(resume_file)
         with open(resume_file, "rb") as f:
-            _resume = pickle.load(f).__dict__
+            _resume = load_data_from_pkl(f).__dict__
 
     strategy = STRATEGIES["automixedprecision"](
         model=wrapped_model,

neural_compressor/quantization.py

@@ -27,7 +27,7 @@
 from .model import Model
 from .strategy import STRATEGIES
 from .utils import logger
-from .utils.utility import dump_class_attrs, secure_check_eval_func, time_limit
+from .utils.utility import dump_class_attrs, load_data_from_pkl, secure_check_eval_func, time_limit
 
 
 def fit(
@@ -182,7 +182,7 @@ def eval_func(model):
     if resume_file:
         assert os.path.exists(resume_file), "The specified resume file {} doesn't exist!".format(resume_file)
         with open(resume_file, "rb") as f:
-            _resume = pickle.load(f).__dict__
+            _resume = load_data_from_pkl(f).__dict__
 
     if eval_func is None and eval_dataloader is None:  # pragma: no cover
         logger.info("Quantize model without tuning!")

neural_compressor/training.py

@@ -32,7 +32,7 @@
 from .metric import register_customer_metric
 from .model.model import Model
 from .utils import logger
-from .utils.utility import time_limit
+from .utils.utility import load_data_from_pkl, time_limit
 
 
 class CompressionManager:
@@ -312,7 +312,7 @@ def eval_func(model):
     if resume_file:
         assert os.path.exists(resume_file), "The specified resume file {} doesn't exist!".format(resume_file)
         with open(resume_file, "rb") as f:
-            _resume = pickle.load(f).__dict__
+            _resume = load_data_from_pkl(f).__dict__
 
     if eval_func is None and eval_dataloader is None:  # pragma: no cover
         logger.info("Quantize model without tuning!")

neural_compressor/utils/utility.py

@@ -638,20 +638,18 @@ class GLOBAL_STATE:
     STATE = MODE.QUANTIZATION
 
 
-def load_data_from_pkl(path, filename):
+def load_data_from_pkl(file_path):
     """Load data from local pkl file.
 
     Args:
-        path: The directory to load data
-        filename: The filename to load
+        file_path: The directory to load data
     """
     try:
-        file_path = os.path.join(path, filename)
         with open(file_path, "rb") as fp:
             data = _safe_pickle_load(fp)
             return data
     except FileExistsError:
-        logging.getLogger("neural_compressor").info("Can not open %s." % path)
+        logging.getLogger("neural_compressor").info("Can not open %s." % file_path)
 
 
 def dump_data_to_local(data, path, filename):