From 1d1610322e19ec91464212a534de95ac5adcdcf0 Mon Sep 17 00:00:00 2001 From: "Cheng, Penghui" Date: Mon, 24 Jun 2024 14:54:06 +0800 Subject: [PATCH] Check neural-compressor version for export function Signed-off-by: Cheng, Penghui --- optimum/intel/neural_compressor/quantization.py | 7 ++++++- optimum/intel/neural_compressor/trainer.py | 6 +++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/optimum/intel/neural_compressor/quantization.py b/optimum/intel/neural_compressor/quantization.py index a345d90a03..90c0c39d68 100644 --- a/optimum/intel/neural_compressor/quantization.py +++ b/optimum/intel/neural_compressor/quantization.py @@ -28,7 +28,6 @@ from neural_compressor.model.onnx_model import ONNXModel from neural_compressor.model.torch_model import IPEXModel, PyTorchModel from neural_compressor.quantization import fit -from neural_compressor.utils.export import torch_to_int8_onnx from packaging.version import parse from torch.utils.data import DataLoader, RandomSampler from transformers import ( @@ -80,6 +79,12 @@ ) +if is_neural_compressor_version("<", "2.6"): + from neural_compressor.experimental.export import torch_to_int8_onnx +else: + from neural_compressor.utils.export import torch_to_int8_onnx + + if is_itrex_available(): if is_itrex_version("<", ITREX_MINIMUM_VERSION): raise ImportError( diff --git a/optimum/intel/neural_compressor/trainer.py b/optimum/intel/neural_compressor/trainer.py index fe17113724..df43b43582 100644 --- a/optimum/intel/neural_compressor/trainer.py +++ b/optimum/intel/neural_compressor/trainer.py @@ -36,7 +36,6 @@ from neural_compressor import training from neural_compressor.compression import DistillationCallbacks from neural_compressor.conf.pythonic_config import _BaseQuantizationConfig -from neural_compressor.utils.export import torch_to_fp32_onnx, torch_to_int8_onnx from packaging import version from torch import nn from torch.utils.data import Dataset, RandomSampler @@ -107,6 +106,11 @@ if TYPE_CHECKING: from optimum.exporters.onnx import OnnxConfig +if is_neural_compressor_version("<", "2.6"): + from neural_compressor.experimental.export import torch_to_fp32_onnx, torch_to_int8_onnx +else: + from neural_compressor.utils.export import torch_to_fp32_onnx, torch_to_int8_onnx + __version__ = "4.22.2"