From 7c3b6454cf7f91b6f3a9a16aa40d7b6ca300646b Mon Sep 17 00:00:00 2001 From: chensuyue Date: Wed, 5 Nov 2025 16:26:26 +0800 Subject: [PATCH 1/6] try safe load Signed-off-by: chensuyue --- neural_compressor/quantization.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/neural_compressor/quantization.py b/neural_compressor/quantization.py index 2f3b2460d98..42297886885 100644 --- a/neural_compressor/quantization.py +++ b/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, secure_check_eval_func, time_limit, _safe_pickle_load 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 = _safe_pickle_load(f).__dict__ if eval_func is None and eval_dataloader is None: # pragma: no cover logger.info("Quantize model without tuning!") From 7ec1618131f9ffb5c4c3bcec4f47e68a5decc90b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 5 Nov 2025 08:30:57 +0000 Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- neural_compressor/quantization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/neural_compressor/quantization.py b/neural_compressor/quantization.py index 42297886885..505fa4bdbfc 100644 --- a/neural_compressor/quantization.py +++ b/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, _safe_pickle_load +from .utils.utility import _safe_pickle_load, dump_class_attrs, secure_check_eval_func, time_limit def fit( From 55c30da208547a656e21b7a2ea4475d61eea9d1b Mon Sep 17 00:00:00 2001 From: chensuyue Date: Fri, 7 Nov 2025 10:41:05 +0800 Subject: [PATCH 3/6] use _safe_pickle_load in source code Signed-off-by: chensuyue --- neural_compressor/data/datasets/dataset.py | 6 +++--- neural_compressor/mix_precision.py | 4 ++-- neural_compressor/training.py | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/neural_compressor/data/datasets/dataset.py b/neural_compressor/data/datasets/dataset.py index 4f97340c8ea..0a169081710 100644 --- a/neural_compressor/data/datasets/dataset.py +++ b/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)} diff --git a/neural_compressor/mix_precision.py b/neural_compressor/mix_precision.py index 84a47025a65..c191ea71d77 100644 --- a/neural_compressor/mix_precision.py +++ b/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, diff --git a/neural_compressor/training.py b/neural_compressor/training.py index 1f1ecfb53fc..c3d9acae3c6 100644 --- a/neural_compressor/training.py +++ b/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!") From 9a87c1e53d3a0fdc3bb58a1e43a1f431477ca5b5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 7 Nov 2025 02:43:18 +0000 Subject: [PATCH 4/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- neural_compressor/data/datasets/dataset.py | 2 +- neural_compressor/mix_precision.py | 2 +- neural_compressor/training.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/neural_compressor/data/datasets/dataset.py b/neural_compressor/data/datasets/dataset.py index 0a169081710..ea1bd8bd3fb 100644 --- a/neural_compressor/data/datasets/dataset.py +++ b/neural_compressor/data/datasets/dataset.py @@ -21,7 +21,7 @@ from PIL import Image -from neural_compressor.utils.utility import LazyImport, singleton, _safe_pickle_load +from neural_compressor.utils.utility import LazyImport, _safe_pickle_load, singleton torch = LazyImport("torch") torchvision = LazyImport("torchvision") diff --git a/neural_compressor/mix_precision.py b/neural_compressor/mix_precision.py index c191ea71d77..685990fd719 100644 --- a/neural_compressor/mix_precision.py +++ b/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, _safe_pickle_load +from .utils.utility import CpuInfo, _safe_pickle_load, secure_check_eval_func, time_limit @alias_param("conf", param_alias="config") diff --git a/neural_compressor/training.py b/neural_compressor/training.py index c3d9acae3c6..5b221fa7e11 100644 --- a/neural_compressor/training.py +++ b/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, _safe_pickle_load +from .utils.utility import _safe_pickle_load, time_limit class CompressionManager: From fcd55ae61d3126999d95644faae1a18e408d98c7 Mon Sep 17 00:00:00 2001 From: chensuyue Date: Fri, 7 Nov 2025 11:13:10 +0800 Subject: [PATCH 5/6] replace with public function Signed-off-by: chensuyue --- neural_compressor/adaptor/tensorflow.py | 2 +- neural_compressor/data/datasets/dataset.py | 6 +++--- neural_compressor/mix_precision.py | 4 ++-- neural_compressor/quantization.py | 4 ++-- neural_compressor/training.py | 4 ++-- neural_compressor/utils/utility.py | 8 +++----- 6 files changed, 13 insertions(+), 15 deletions(-) diff --git a/neural_compressor/adaptor/tensorflow.py b/neural_compressor/adaptor/tensorflow.py index 91ecbb220f0..e03260de409 100644 --- a/neural_compressor/adaptor/tensorflow.py +++ b/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 = {} diff --git a/neural_compressor/data/datasets/dataset.py b/neural_compressor/data/datasets/dataset.py index ea1bd8bd3fb..4f97340c8ea 100644 --- a/neural_compressor/data/datasets/dataset.py +++ b/neural_compressor/data/datasets/dataset.py @@ -21,7 +21,7 @@ from PIL import Image -from neural_compressor.utils.utility import LazyImport, _safe_pickle_load, singleton +from neural_compressor.utils.utility import LazyImport, singleton 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 = _safe_pickle_load(f, encoding="latin1") + entry = 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 = _safe_pickle_load(infile, encoding="latin1") + data = pickle.load(infile, encoding="latin1") self.classes = data[self.meta["key"]] self.class_to_idx = {_class: i for i, _class in enumerate(self.classes)} diff --git a/neural_compressor/mix_precision.py b/neural_compressor/mix_precision.py index 685990fd719..bbeae0aa681 100644 --- a/neural_compressor/mix_precision.py +++ b/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, _safe_pickle_load, secure_check_eval_func, time_limit +from .utils.utility import CpuInfo, secure_check_eval_func, time_limit, load_data_from_pkl @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 = _safe_pickle_load(f).__dict__ + _resume = load_data_from_pkl(f).__dict__ strategy = STRATEGIES["automixedprecision"]( model=wrapped_model, diff --git a/neural_compressor/quantization.py b/neural_compressor/quantization.py index 505fa4bdbfc..0efde0075b2 100644 --- a/neural_compressor/quantization.py +++ b/neural_compressor/quantization.py @@ -27,7 +27,7 @@ from .model import Model from .strategy import STRATEGIES from .utils import logger -from .utils.utility import _safe_pickle_load, dump_class_attrs, secure_check_eval_func, time_limit +from .utils.utility import dump_class_attrs, secure_check_eval_func, time_limit, load_data_from_pkl 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 = _safe_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!") diff --git a/neural_compressor/training.py b/neural_compressor/training.py index 5b221fa7e11..704adef4a2e 100644 --- a/neural_compressor/training.py +++ b/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 _safe_pickle_load, 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 = _safe_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!") diff --git a/neural_compressor/utils/utility.py b/neural_compressor/utils/utility.py index bc3519678b9..2f1a7836a00 100644 --- a/neural_compressor/utils/utility.py +++ b/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): From 3bc6fbd6873abb5745c3c950ff452b2418570d49 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 7 Nov 2025 03:17:25 +0000 Subject: [PATCH 6/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- neural_compressor/mix_precision.py | 2 +- neural_compressor/quantization.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/neural_compressor/mix_precision.py b/neural_compressor/mix_precision.py index bbeae0aa681..ef50ff342b6 100644 --- a/neural_compressor/mix_precision.py +++ b/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, load_data_from_pkl +from .utils.utility import CpuInfo, load_data_from_pkl, secure_check_eval_func, time_limit @alias_param("conf", param_alias="config") diff --git a/neural_compressor/quantization.py b/neural_compressor/quantization.py index 0efde0075b2..947dd0dc3e3 100644 --- a/neural_compressor/quantization.py +++ b/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, load_data_from_pkl +from .utils.utility import dump_class_attrs, load_data_from_pkl, secure_check_eval_func, time_limit def fit(