use _safe_pickle_load in source code

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

Author: chensuyue

neural_compressor/data/datasets/dataset.py

@@ -21,7 +21,7 @@
 
 from PIL import Image
 
-from neural_compressor.utils.utility import LazyImport, singleton
+from neural_compressor.utils.utility import LazyImport, singleton, _safe_pickle_load
 
 torch = LazyImport("torch")
 torchvision = LazyImport("torchvision")
@@ -415,7 +415,7 @@ def __init__(self, root, train=False, transform=None, filter=None, download=True
         for file_name, checksum in downloaded_list:
             file_path = os.path.join(self.root, file_name)
             with open(file_path, "rb") as f:
-                entry = pickle.load(f, encoding="latin1")
+                entry = _safe_pickle_load(f, encoding="latin1")
                 self.data.append(entry["data"])
                 if "labels" in entry:
                     self.targets.extend(entry["labels"])
@@ -435,7 +435,7 @@ def load_meta(self):  # pragma: no cover
                 "Dataset metadata file not found or corrupted." + " You can use download=True to download it"
             )
         with open(path, "rb") as infile:
-            data = pickle.load(infile, encoding="latin1")
+            data = _safe_pickle_load(infile, encoding="latin1")
             self.classes = data[self.meta["key"]]
         self.class_to_idx = {_class: i for i, _class in enumerate(self.classes)}
 

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, secure_check_eval_func, time_limit, _safe_pickle_load
 
 
 @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 = _safe_pickle_load(f).__dict__
 
     strategy = STRATEGIES["automixedprecision"](
         model=wrapped_model,

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 time_limit, _safe_pickle_load
 
 
 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 = _safe_pickle_load(f).__dict__
 
     if eval_func is None and eval_dataloader is None:  # pragma: no cover
         logger.info("Quantize model without tuning!")