From ef8e20ef24f230c46aa48441df3948c32803fb13 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Fri, 10 Jan 2025 11:06:17 +0100 Subject: [PATCH 01/22] add type checking --- optimum/utils/normalized_config.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/optimum/utils/normalized_config.py b/optimum/utils/normalized_config.py index 053417b20b..8b257ba334 100644 --- a/optimum/utils/normalized_config.py +++ b/optimum/utils/normalized_config.py @@ -15,9 +15,11 @@ """Normalization configuration classes.""" import functools -from typing import Callable, Dict, Type, Union +from typing import TYPE_CHECKING, Callable, Dict, Type, Union -from transformers import PretrainedConfig + +if TYPE_CHECKING: + from transformers import PretrainedConfig class NormalizedConfig: @@ -29,7 +31,7 @@ class NormalizedConfig: The config to normalize. """ - def __init__(self, config: Union[PretrainedConfig, Dict], allow_new: bool = False, **kwargs): + def __init__(self, config: Union["PretrainedConfig", Dict], allow_new: bool = False, **kwargs): self.config = config for key, value in kwargs.items(): if allow_new or hasattr(self, key.upper()): @@ -40,7 +42,7 @@ def __init__(self, config: Union[PretrainedConfig, Dict], allow_new: bool = Fals ) @classmethod - def with_args(cls, allow_new: bool = False, **kwargs) -> Callable[[PretrainedConfig], "NormalizedConfig"]: + def with_args(cls, allow_new: bool = False, **kwargs) -> Callable[["PretrainedConfig"], "NormalizedConfig"]: return functools.partial(cls, allow_new=allow_new, **kwargs) def __getattr__(self, attr_name): From 1493f70eb97a83ccf80ee49eccda07fd44725d0f Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Wed, 15 Jan 2025 16:16:26 +0100 Subject: [PATCH 02/22] tflite config exporters --- optimum/exporters/tflite/base.py | 72 ++++---------------------------- 1 file changed, 8 insertions(+), 64 deletions(-) diff --git a/optimum/exporters/tflite/base.py b/optimum/exporters/tflite/base.py index 3df230c33b..e223d398dc 100644 --- a/optimum/exporters/tflite/base.py +++ b/optimum/exporters/tflite/base.py @@ -27,7 +27,7 @@ if is_tf_available(): import tensorflow as tf -from ..base import ExportConfig +from ..base import ExportersConfig if TYPE_CHECKING: @@ -115,7 +115,7 @@ def __post_init__(self): self.approach = QuantizationApproach(self.approach) -class TFLiteConfig(ExportConfig, ABC): +class TFLiteConfig(ExportersConfig, ABC): """ Base class for TFLite exportable model describing metadata on how to export the model through the TFLite format. @@ -192,11 +192,11 @@ def __init__( point_batch_size: Optional[int] = None, nb_points_per_image: Optional[int] = None, ): - self._config = config - self._normalized_config = self.NORMALIZED_CONFIG_CLASS(self._config) - self.mandatory_axes = () - self.task = task - self._axes: Dict[str, int] = {} + + super().__init__(config=config, task=task, int_dtype="int64", float_dtype="fp32") + + # self.mandatory_axes = () + # self._axes: Dict[str, int] = {} # To avoid using **kwargs. axes_values = { @@ -266,65 +266,9 @@ def _create_dummy_input_generator_classes(self) -> List["DummyInputGenerator"]: self._validate_mandatory_axes() return [cls_(self.task, self._normalized_config, **self._axes) for cls_ in self.DUMMY_INPUT_GENERATOR_CLASSES] - @property - def values_override(self) -> Optional[Dict[str, Any]]: - """ - Dictionary of keys to override in the model's config before exporting. - - Returns: - `Optional[Dict[str, Any]]`: A dictionary specifying the configuration items to override. - """ - if hasattr(self._config, "use_cache"): - return {"use_cache": False} - - return None - - @property - @abstractmethod - def inputs(self) -> List[str]: - """ - List containing the names of the inputs the exported model should take. - - Returns: - `List[str]`: A list of input names. - """ - raise NotImplementedError() - - @property - def outputs(self) -> List[str]: - """ - List containing the names of the outputs the exported model should have. - - Returns: - `List[str]`: A list of output names. - """ - return self._TASK_TO_COMMON_OUTPUTS[self.task] - def generate_dummy_inputs(self) -> Dict[str, "tf.Tensor"]: - """ - Generates dummy inputs that the exported model should be able to process. - This method is actually used to determine the input specs that are needed for the export. - - Returns: - `Dict[str, tf.Tensor]`: A dictionary mapping input names to dummy tensors. - """ - dummy_inputs_generators = self._create_dummy_input_generator_classes() - dummy_inputs = {} - - for input_name in self.inputs: - input_was_inserted = False - for dummy_input_gen in dummy_inputs_generators: - if dummy_input_gen.supports_input(input_name): - dummy_inputs[input_name] = dummy_input_gen.generate(input_name, framework="tf") - input_was_inserted = True - break - if not input_was_inserted: - raise RuntimeError( - f'Could not generate dummy inputs for "{input_name}". Try adding a proper dummy input generator ' - "to the model TFLite config." - ) + return super().generate_dummy_inputs(framework="tf") - return dummy_inputs @property def inputs_specs(self) -> List["TensorSpec"]: From 58b906f60f885ecbe5d2c5bad4bead682ea67fab Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Wed, 15 Jan 2025 17:18:17 +0100 Subject: [PATCH 03/22] add exporters config --- optimum/exporters/base.py | 286 +++++++++++++++++++++++++++++++++ optimum/exporters/onnx/base.py | 145 +---------------- 2 files changed, 293 insertions(+), 138 deletions(-) diff --git a/optimum/exporters/base.py b/optimum/exporters/base.py index 17e1265e74..2080e9890f 100644 --- a/optimum/exporters/base.py +++ b/optimum/exporters/base.py @@ -17,5 +17,291 @@ from abc import ABC + +import copy +import enum +import gc +import inspect +import itertools +import os +import re +from abc import ABC, abstractmethod +from collections import OrderedDict +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union + +import numpy as np +from transformers.utils import is_accelerate_available, is_torch_available + + +if is_torch_available(): + import torch.nn as nn + +from .utils import ( + DEFAULT_DUMMY_SHAPES, + DummyInputGenerator, + logging, +) +from .utils import TORCH_MINIMUM_VERSION as GLOBAL_MIN_TORCH_VERSION +from .utils import TRANSFORMERS_MINIMUM_VERSION as GLOBAL_MIN_TRANSFORMERS_VERSION +from .utils.doc import add_dynamic_docstring +from .utils.import_utils import is_torch_version, is_transformers_version +from .error_utils import MissingMandatoryAxisDimension + +# from .model_patcher import ModelPatcher + +if TYPE_CHECKING: + from transformers import PretrainedConfig, PreTrainedModel, TFPreTrainedModel + + from .model_patcher import PatchingSpec + +logger = logging.get_logger(__name__) + + + +GENERATE_DUMMY_DOCSTRING = r""" + Generates the dummy inputs necessary for tracing the model. If not explicitely specified, default input shapes are used. + + Args: + framework (`str`, defaults to `"pt"`): + The framework for which to create the dummy inputs. + batch_size (`int`, defaults to {batch_size}): + The batch size to use in the dummy inputs. + sequence_length (`int`, defaults to {sequence_length}): + The sequence length to use in the dummy inputs. + num_choices (`int`, defaults to {num_choices}): + The number of candidate answers provided for multiple choice task. + image_width (`int`, defaults to {width}): + The width to use in the dummy inputs for vision tasks. + image_height (`int`, defaults to {height}): + The height to use in the dummy inputs for vision tasks. + num_channels (`int`, defaults to {num_channels}): + The number of channels to use in the dummpy inputs for vision tasks. + feature_size (`int`, defaults to {feature_size}): + The number of features to use in the dummpy inputs for audio tasks in case it is not raw audio. + This is for example the number of STFT bins or MEL bins. + nb_max_frames (`int`, defaults to {nb_max_frames}): + The number of frames to use in the dummpy inputs for audio tasks in case the input is not raw audio. + audio_sequence_length (`int`, defaults to {audio_sequence_length}): + The number of frames to use in the dummpy inputs for audio tasks in case the input is raw audio. + + Returns: + `Dict`: A dictionary mapping the input names to dummy tensors in the proper framework format. +""" + + + +# TODO: Remove class ExportConfig(ABC): pass + + + +class ExportersConfig(ABC): + """ + Base class describing metadata on how to export the model through the ONNX format. + + Class attributes: + + - NORMALIZED_CONFIG_CLASS (`Type`) -- A class derived from [`~optimum.utils.NormalizedConfig`] specifying how to + normalize the model config. + - DUMMY_INPUT_GENERATOR_CLASSES (`Tuple[Type]`) -- A tuple of classes derived from + [`~optimum.utils.DummyInputGenerator`] specifying how to create dummy inputs. + - ATOL_FOR_VALIDATION (`Union[float, Dict[str, float]]`) -- A float or a dictionary mapping task names to float, + where the float values represent the absolute tolerance value to use during model conversion validation. + - MIN_TORCH_VERSION (`packaging.version.Version`, defaults to [`~optimum.exporters.utils.TORCH_MINIMUM_VERSION`]) -- The + minimum torch version supporting the export of the model. + - MIN_TRANSFORMERS_VERSION (`packaging.version.Version`, defaults to + [`~optimum.exporters.utils.TRANSFORMERS_MINIMUM_VERSION`] -- The minimum transformers version supporting the + export of the model. Not always up-to-date or accurate. This is more for internal use. + - PATCHING_SPECS (`Optional[List[PatchingSpec]]`, defaults to `None`) -- Specify which operators / modules should be + patched before performing the export, and how. This is useful when some operator is not supported for instance. + + Args: + config (`transformers.PretrainedConfig`): + The model configuration. + task (`str`, defaults to `"feature-extraction"`): + The task the model should be exported for. + int_dtype (`str`, defaults to `"int64"`): + The data type of integer tensors, could be ["int64", "int32", "int8"], default to "int64". + float_dtype (`str`, defaults to `"fp32"`): + The data type of float tensors, could be ["fp32", "fp16", "bf16"], default to "fp32". + """ + + NORMALIZED_CONFIG_CLASS = None + DUMMY_INPUT_GENERATOR_CLASSES = () + ATOL_FOR_VALIDATION: Union[float, Dict[str, float]] = 1e-5 + MIN_TORCH_VERSION = GLOBAL_MIN_TORCH_VERSION + MIN_TRANSFORMERS_VERSION = GLOBAL_MIN_TRANSFORMERS_VERSION + PATCHING_SPECS: Optional[List["PatchingSpec"]] = None + _TASK_TO_COMMON_OUTPUTS = { + "audio-classification": ["logits"], + "audio-frame-classification": ["logits"], + "automatic-speech-recognition": ["logits"], + "audio-xvector": ["logits"], # for onnx : ["logits", "embeddings"] + "depth-estimation": ["predicted_depth"], + "document-question-answering": ["logits"], + "feature-extraction": ["last_hidden_state"], # for neuron : ["last_hidden_state", "pooler_output"] + "fill-mask": ["logits"], + "image-classification": ["logits"], + "image-segmentation": ["logits"], # for tflite : ["logits", "pred_boxes", "pred_masks"] + "image-to-text": ["logits"], + "image-to-image": ["reconstruction"], + "mask-generation": ["logits"], + "masked-im": ["logits"], # for onnx : ["reconstruction"] + "multiple-choice": ["logits"], + "object-detection": ["logits", "pred_boxes"], + "question-answering": ["start_logits", "end_logits"], + "semantic-segmentation": ["logits"], + "text2text-generation": ["logits"], # for tflite : ["logits", "encoder_last_hidden_state"], + "text-classification": ["logits"], + "text-generation": ["logits"], + "time-series-forecasting": ["prediction_outputs"], + "token-classification": ["logits"], + "visual-question-answering": ["logits"], + "zero-shot-image-classification": ["logits_per_image", "logits_per_text", "text_embeds", "image_embeds"], + "zero-shot-object-detection": ["logits", "pred_boxes", "text_embeds", "image_embeds"], + } + + def __init__( + self, + config: "PretrainedConfig", + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + ): + self.task = task + self._config = config + self._normalized_config = self.NORMALIZED_CONFIG_CLASS(self._config) + self.int_dtype = int_dtype + self.float_dtype = float_dtype + + self.mandatory_axes = () + self._axes: Dict[str, int] = {} + + + def _create_dummy_input_generator_classes(self, **kwargs) -> List[DummyInputGenerator]: + """ + Instantiates the dummy input generators from `self.DUMMY_INPUT_GENERATOR_CLASSES`. + Each dummy input generator is independent, so this method instantiates the first generator, and + forces the other generators to use the same batch size, meaning they will all produce inputs of the same batch + size. Override this method for custom behavior. + """ + # self._validate_mandatory_axes() + return [cls_(self.task, self._normalized_config, **kwargs) for cls_ in self.DUMMY_INPUT_GENERATOR_CLASSES] + + + @property + @abstractmethod + def inputs(self) -> Dict[str, Dict[int, str]]: + """ + Dict containing the axis definition of the input tensors to provide to the model. + + Returns: + `Dict[str, Dict[int, str]]`: A mapping of each input name to a mapping of axis position to the axes symbolic name. + """ + raise NotImplementedError() + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + """ + Dict containing the axis definition of the output tensors to provide to the model. + + Returns: + `Dict[str, Dict[int, str]]`: A mapping of each output name to a mapping of axis position to the axes symbolic name. + """ + common_outputs = self._TASK_TO_COMMON_OUTPUTS[self.task] + return copy.deepcopy(common_outputs) + + + @property + def values_override(self) -> Optional[Dict[str, Any]]: + """ + Dictionary of keys to override in the model's config before exporting. + + Returns: + `Optional[Dict[str, Any]]`: A dictionary specifying the configuration items to override. + """ + if hasattr(self._config, "use_cache"): + return {"use_cache": False} + + return None + + @property + def is_transformers_support_available(self) -> bool: + """ + Whether the installed version of Transformers allows for the ONNX export. + + Returns: + `bool`: Whether the install version of Transformers is compatible with the model. + + """ + return is_transformers_version(">=", self.MIN_TRANSFORMERS_VERSION.base_version) + + @property + def is_torch_support_available(self) -> bool: + """ + Whether the installed version of PyTorch allows for the ONNX export. + + Returns: + `bool`: Whether the installed version of PyTorch is compatible with the model. + """ + if is_torch_available(): + return is_torch_version(">=", self.MIN_TORCH_VERSION.base_version) + + return False + + + @add_dynamic_docstring(text=GENERATE_DUMMY_DOCSTRING, dynamic_elements=DEFAULT_DUMMY_SHAPES) + def generate_dummy_inputs(self, framework: str = "pt", **kwargs) -> Dict: + + """ + Generates dummy inputs that the exported model should be able to process. + This method is actually used to determine the input specs that are needed for the export. + + Returns: + `Dict[str, [tf.Tensor, torch.Tensor]]`: A dictionary mapping input names to dummy tensors. + """ + + + dummy_inputs_generators = self._create_dummy_input_generator_classes(**kwargs) + dummy_inputs = {} + for input_name in self.inputs: + input_was_inserted = False + for dummy_input_gen in dummy_inputs_generators: + if dummy_input_gen.supports_input(input_name): + dummy_inputs[input_name] = dummy_input_gen.generate( + input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype + ) + input_was_inserted = True + break + if not input_was_inserted: + raise RuntimeError( + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to ' + "the model exporters config." + ) + return dummy_inputs + + @classmethod + def flatten_inputs(cls, inputs: Dict[str, Any]) -> Dict[str, Any]: + """ + Flatten nested structure in dummy inputs, e.g `addition_embed_type` of unet model. + """ + flatten = {} + for name, value in inputs.items(): + if isinstance(value, dict): + for sub_name, sub_value in value.items(): + flatten[sub_name] = sub_value + else: + flatten[name] = value + return flatten + + ############################################################################################################################################################ + + # def patch_model_for_export( + # self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None + # ) -> ModelPatcher: + # return ModelPatcher(self, model, model_kwargs=model_kwargs) + + + ############################################################################################################################################################ diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py index e6618568c0..b8498e8cfe 100644 --- a/optimum/exporters/onnx/base.py +++ b/optimum/exporters/onnx/base.py @@ -50,7 +50,7 @@ is_torch_version, is_transformers_version, ) -from ..base import ExportConfig +from ..base import ExportConfig, ExportersConfig from .constants import ONNX_DECODER_MERGED_NAME, ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME from .model_patcher import ModelPatcher, Seq2SeqModelPatcher @@ -102,48 +102,12 @@ """ -class OnnxConfig(ExportConfig, ABC): - """ - Base class for ONNX exportable model describing metadata on how to export the model through the ONNX format. - - Class attributes: - - - NORMALIZED_CONFIG_CLASS (`Type`) -- A class derived from [`~optimum.utils.NormalizedConfig`] specifying how to - normalize the model config. - - DUMMY_INPUT_GENERATOR_CLASSES (`Tuple[Type]`) -- A tuple of classes derived from - [`~optimum.utils.DummyInputGenerator`] specifying how to create dummy inputs. - - ATOL_FOR_VALIDATION (`Union[float, Dict[str, float]]`) -- A float or a dictionary mapping task names to float, - where the float values represent the absolute tolerance value to use during model conversion validation. - - DEFAULT_ONNX_OPSET (`int`, defaults to 11) -- The default ONNX opset to use for the ONNX export. - - MIN_TORCH_VERSION (`packaging.version.Version`, defaults to [`~optimum.exporters.onnx.utils.TORCH_MINIMUM_VERSION`]) -- The - minimum torch version supporting the export of the model to ONNX. - - MIN_TRANSFORMERS_VERSION (`packaging.version.Version`, defaults to - [`~optimum.exporters.onnx.utils.TRANSFORMERS_MINIMUM_VERSION`] -- The minimum transformers version supporting the - export of the model to ONNX. Not always up-to-date or accurate. This is more for internal use. - - PATCHING_SPECS (`Optional[List[PatchingSpec]]`, defaults to `None`) -- Specify which operators / modules should be - patched before performing the export, and how. This is useful when some operator is not supported in ONNX for - instance. - - Args: - config (`transformers.PretrainedConfig`): - The model configuration. - task (`str`, defaults to `"feature-extraction"`): - The task the model should be exported for. - int_dtype (`str`, defaults to `"int64"`): - The data type of integer tensors, could be ["int64", "int32", "int8"], default to "int64". - float_dtype (`str`, defaults to `"fp32"`): - The data type of float tensors, could be ["fp32", "fp16", "bf16"], default to "fp32". - """ +class OnnxConfig(ExportersConfig): - NORMALIZED_CONFIG_CLASS = None - DUMMY_INPUT_GENERATOR_CLASSES = () DEFAULT_ONNX_OPSET = 11 - ATOL_FOR_VALIDATION: Union[float, Dict[str, float]] = 1e-5 - MIN_TORCH_VERSION = GLOBAL_MIN_TORCH_VERSION - MIN_TRANSFORMERS_VERSION = GLOBAL_MIN_TRANSFORMERS_VERSION - PATCHING_SPECS: Optional[List["PatchingSpec"]] = None VARIANTS = {"default": "The default ONNX variant."} DEFAULT_VARIANT = "default" + _TASK_TO_COMMON_OUTPUTS = { "audio-classification": OrderedDict({"logits": {0: "batch_size"}}), "audio-frame-classification": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), @@ -210,53 +174,12 @@ def __init__( float_dtype: str = "fp32", legacy: bool = False, ): - self.task = task - self.int_dtype = int_dtype - self.float_dtype = float_dtype + super().__init__(config=config, task=task, int_dtype=int_dtype, float_dtype=float_dtype) - self._config = config - self._preprocessors = preprocessors - self._normalized_config = self.NORMALIZED_CONFIG_CLASS(self._config) self.variant = "default" + self._preprocessors = preprocessors self.legacy = legacy - def _create_dummy_input_generator_classes(self, **kwargs) -> List[DummyInputGenerator]: - """ - Instantiates the dummy input generators from `self.DUMMY_INPUT_GENERATOR_CLASSES`. - Each dummy input generator is independent, so this method instantiates the first generator, and - forces the other generators to use the same batch size, meaning they will all produce inputs of the same batch - size. Override this method for custom behavior. - """ - first_inputs_gen = self.DUMMY_INPUT_GENERATOR_CLASSES[0](self.task, self._normalized_config, **kwargs) - dummy_inputs_generators = [ - cls_(self.task, self._normalized_config, **kwargs) for cls_ in self.DUMMY_INPUT_GENERATOR_CLASSES[1:] - ] - dummy_inputs_generators.insert(0, first_inputs_gen) - - return dummy_inputs_generators - - @property - @abstractmethod - def inputs(self) -> Dict[str, Dict[int, str]]: - """ - Dict containing the axis definition of the input tensors to provide to the model. - - Returns: - `Dict[str, Dict[int, str]]`: A mapping of each input name to a mapping of axis position to the axes symbolic name. - """ - raise NotImplementedError() - - @property - def outputs(self) -> Dict[str, Dict[int, str]]: - """ - Dict containing the axis definition of the output tensors to provide to the model. - - Returns: - `Dict[str, Dict[int, str]]`: A mapping of each output name to a mapping of axis position to the axes symbolic name. - """ - common_outputs = self._TASK_TO_COMMON_OUTPUTS[self.task] - return copy.deepcopy(common_outputs) - @property def variant(self) -> str: """ @@ -358,43 +281,7 @@ def patch_model_for_export( self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None ) -> ModelPatcher: return ModelPatcher(self, model, model_kwargs=model_kwargs) - - @property - def values_override(self) -> Optional[Dict[str, Any]]: - """ - Dictionary of keys to override in the model's config before exporting. - - Returns: - `Optional[Dict[str, Any]]`: A dictionary specifying the configuration items to override. - """ - if hasattr(self._config, "use_cache"): - return {"use_cache": False} - - return None - - @property - def is_transformers_support_available(self) -> bool: - """ - Whether the installed version of Transformers allows for the ONNX export. - - Returns: - `bool`: Whether the install version of Transformers is compatible with the model. - - """ - return is_transformers_version(">=", self.MIN_TRANSFORMERS_VERSION.base_version) - - @property - def is_torch_support_available(self) -> bool: - """ - Whether the installed version of PyTorch allows for the ONNX export. - - Returns: - `bool`: Whether the installed version of PyTorch is compatible with the model. - """ - if is_torch_available(): - return is_torch_version(">=", self.MIN_TORCH_VERSION.base_version) - - return False + @property def torch_to_onnx_input_map(self) -> Dict[str, str]: @@ -461,27 +348,9 @@ def ordered_inputs(self, model: Union["PreTrainedModel", "TFPreTrainedModel"]) - ordered_inputs[name] = dynamic_axes return ordered_inputs - @add_dynamic_docstring(text=GENERATE_DUMMY_DOCSTRING, dynamic_elements=DEFAULT_DUMMY_SHAPES) - def generate_dummy_inputs(self, framework: str = "pt", **kwargs) -> Dict: - dummy_inputs_generators = self._create_dummy_input_generator_classes(**kwargs) - dummy_inputs = {} - for input_name in self.inputs: - input_was_inserted = False - for dummy_input_gen in dummy_inputs_generators: - if dummy_input_gen.supports_input(input_name): - dummy_inputs[input_name] = dummy_input_gen.generate( - input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype - ) - input_was_inserted = True - break - if not input_was_inserted: - raise RuntimeError( - f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to ' - "the model ONNX config." - ) - return dummy_inputs + # TODO: use instead flatten_inputs and remove @classmethod def flatten_output_collection_property(cls, name: str, field: Iterable[Any]) -> Dict[str, Any]: """ From 46f1c263c698c94df9b914b0ee90fe9041b9f449 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Wed, 15 Jan 2025 17:19:08 +0100 Subject: [PATCH 04/22] fix style --- optimum/exporters/base.py | 39 +++++++------------------------- optimum/exporters/onnx/base.py | 12 ++-------- optimum/exporters/tflite/base.py | 4 +--- optimum/exporters/utils.py | 2 +- 4 files changed, 12 insertions(+), 45 deletions(-) diff --git a/optimum/exporters/base.py b/optimum/exporters/base.py index 2080e9890f..7260343969 100644 --- a/optimum/exporters/base.py +++ b/optimum/exporters/base.py @@ -14,28 +14,15 @@ # limitations under the License. """Base exporters config.""" -from abc import ABC - - - import copy -import enum -import gc -import inspect -import itertools -import os -import re from abc import ABC, abstractmethod -from collections import OrderedDict -from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union -import numpy as np -from transformers.utils import is_accelerate_available, is_torch_available +from transformers.utils import is_torch_available if is_torch_available(): - import torch.nn as nn + pass from .utils import ( DEFAULT_DUMMY_SHAPES, @@ -46,19 +33,18 @@ from .utils import TRANSFORMERS_MINIMUM_VERSION as GLOBAL_MIN_TRANSFORMERS_VERSION from .utils.doc import add_dynamic_docstring from .utils.import_utils import is_torch_version, is_transformers_version -from .error_utils import MissingMandatoryAxisDimension + # from .model_patcher import ModelPatcher if TYPE_CHECKING: - from transformers import PretrainedConfig, PreTrainedModel, TFPreTrainedModel + from transformers import PretrainedConfig from .model_patcher import PatchingSpec logger = logging.get_logger(__name__) - GENERATE_DUMMY_DOCSTRING = r""" Generates the dummy inputs necessary for tracing the model. If not explicitely specified, default input shapes are used. @@ -90,13 +76,11 @@ """ - # TODO: Remove class ExportConfig(ABC): pass - class ExportersConfig(ABC): """ Base class describing metadata on how to export the model through the ONNX format. @@ -141,19 +125,19 @@ class ExportersConfig(ABC): "audio-xvector": ["logits"], # for onnx : ["logits", "embeddings"] "depth-estimation": ["predicted_depth"], "document-question-answering": ["logits"], - "feature-extraction": ["last_hidden_state"], # for neuron : ["last_hidden_state", "pooler_output"] + "feature-extraction": ["last_hidden_state"], # for neuron : ["last_hidden_state", "pooler_output"] "fill-mask": ["logits"], "image-classification": ["logits"], "image-segmentation": ["logits"], # for tflite : ["logits", "pred_boxes", "pred_masks"] "image-to-text": ["logits"], "image-to-image": ["reconstruction"], "mask-generation": ["logits"], - "masked-im": ["logits"], # for onnx : ["reconstruction"] + "masked-im": ["logits"], # for onnx : ["reconstruction"] "multiple-choice": ["logits"], "object-detection": ["logits", "pred_boxes"], "question-answering": ["start_logits", "end_logits"], "semantic-segmentation": ["logits"], - "text2text-generation": ["logits"], # for tflite : ["logits", "encoder_last_hidden_state"], + "text2text-generation": ["logits"], # for tflite : ["logits", "encoder_last_hidden_state"], "text-classification": ["logits"], "text-generation": ["logits"], "time-series-forecasting": ["prediction_outputs"], @@ -179,7 +163,6 @@ def __init__( self.mandatory_axes = () self._axes: Dict[str, int] = {} - def _create_dummy_input_generator_classes(self, **kwargs) -> List[DummyInputGenerator]: """ Instantiates the dummy input generators from `self.DUMMY_INPUT_GENERATOR_CLASSES`. @@ -190,7 +173,6 @@ def _create_dummy_input_generator_classes(self, **kwargs) -> List[DummyInputGene # self._validate_mandatory_axes() return [cls_(self.task, self._normalized_config, **kwargs) for cls_ in self.DUMMY_INPUT_GENERATOR_CLASSES] - @property @abstractmethod def inputs(self) -> Dict[str, Dict[int, str]]: @@ -213,7 +195,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: common_outputs = self._TASK_TO_COMMON_OUTPUTS[self.task] return copy.deepcopy(common_outputs) - @property def values_override(self) -> Optional[Dict[str, Any]]: """ @@ -251,10 +232,8 @@ def is_torch_support_available(self) -> bool: return False - @add_dynamic_docstring(text=GENERATE_DUMMY_DOCSTRING, dynamic_elements=DEFAULT_DUMMY_SHAPES) def generate_dummy_inputs(self, framework: str = "pt", **kwargs) -> Dict: - """ Generates dummy inputs that the exported model should be able to process. This method is actually used to determine the input specs that are needed for the export. @@ -262,7 +241,6 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs) -> Dict: Returns: `Dict[str, [tf.Tensor, torch.Tensor]]`: A dictionary mapping input names to dummy tensors. """ - dummy_inputs_generators = self._create_dummy_input_generator_classes(**kwargs) dummy_inputs = {} @@ -303,5 +281,4 @@ def flatten_inputs(cls, inputs: Dict[str, Any]) -> Dict[str, Any]: # ) -> ModelPatcher: # return ModelPatcher(self, model, model_kwargs=model_kwargs) - ############################################################################################################################################################ diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py index b8498e8cfe..00e126c2f2 100644 --- a/optimum/exporters/onnx/base.py +++ b/optimum/exporters/onnx/base.py @@ -21,7 +21,7 @@ import itertools import os import re -from abc import ABC, abstractmethod +from abc import ABC from collections import OrderedDict from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union @@ -41,16 +41,13 @@ is_diffusers_available, logging, ) -from ...utils import TORCH_MINIMUM_VERSION as GLOBAL_MIN_TORCH_VERSION -from ...utils import TRANSFORMERS_MINIMUM_VERSION as GLOBAL_MIN_TRANSFORMERS_VERSION from ...utils.doc import add_dynamic_docstring from ...utils.import_utils import ( is_onnx_available, is_onnxruntime_available, - is_torch_version, is_transformers_version, ) -from ..base import ExportConfig, ExportersConfig +from ..base import ExportersConfig from .constants import ONNX_DECODER_MERGED_NAME, ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME from .model_patcher import ModelPatcher, Seq2SeqModelPatcher @@ -66,7 +63,6 @@ if is_diffusers_available(): from diffusers import ModelMixin - from .model_patcher import PatchingSpec logger = logging.get_logger(__name__) @@ -103,7 +99,6 @@ class OnnxConfig(ExportersConfig): - DEFAULT_ONNX_OPSET = 11 VARIANTS = {"default": "The default ONNX variant."} DEFAULT_VARIANT = "default" @@ -281,7 +276,6 @@ def patch_model_for_export( self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None ) -> ModelPatcher: return ModelPatcher(self, model, model_kwargs=model_kwargs) - @property def torch_to_onnx_input_map(self) -> Dict[str, str]: @@ -348,8 +342,6 @@ def ordered_inputs(self, model: Union["PreTrainedModel", "TFPreTrainedModel"]) - ordered_inputs[name] = dynamic_axes return ordered_inputs - - # TODO: use instead flatten_inputs and remove @classmethod def flatten_output_collection_property(cls, name: str, field: Iterable[Any]) -> Dict[str, Any]: diff --git a/optimum/exporters/tflite/base.py b/optimum/exporters/tflite/base.py index e223d398dc..1c1e725163 100644 --- a/optimum/exporters/tflite/base.py +++ b/optimum/exporters/tflite/base.py @@ -14,7 +14,7 @@ # limitations under the License. """TensorFlow Lite configuration base classes.""" -from abc import ABC, abstractmethod +from abc import ABC from ctypes import ArgumentError from dataclasses import dataclass from enum import Enum @@ -192,7 +192,6 @@ def __init__( point_batch_size: Optional[int] = None, nb_points_per_image: Optional[int] = None, ): - super().__init__(config=config, task=task, int_dtype="int64", float_dtype="fp32") # self.mandatory_axes = () @@ -269,7 +268,6 @@ def _create_dummy_input_generator_classes(self) -> List["DummyInputGenerator"]: def generate_dummy_inputs(self) -> Dict[str, "tf.Tensor"]: return super().generate_dummy_inputs(framework="tf") - @property def inputs_specs(self) -> List["TensorSpec"]: """ diff --git a/optimum/exporters/utils.py b/optimum/exporters/utils.py index 58e170ba97..aae07c411d 100644 --- a/optimum/exporters/utils.py +++ b/optimum/exporters/utils.py @@ -545,7 +545,7 @@ def get_speecht5_models_for_export( use_past=use_past, use_past_in_inputs=False, # Irrelevant here. behavior=config._behavior, # Irrelevant here. - preprocessors=config._preprocessors, + # preprocessors=config._preprocessors, is_postnet_and_vocoder=True, legacy=config.legacy, ) From cf9a51d5a66b27d1248a72d1ff928a0101fa0789 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Thu, 16 Jan 2025 16:26:13 +0100 Subject: [PATCH 05/22] fix --- optimum/exporters/base.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/optimum/exporters/base.py b/optimum/exporters/base.py index 7260343969..990651c3f4 100644 --- a/optimum/exporters/base.py +++ b/optimum/exporters/base.py @@ -24,15 +24,15 @@ if is_torch_available(): pass -from .utils import ( +from ..utils import ( DEFAULT_DUMMY_SHAPES, DummyInputGenerator, logging, ) -from .utils import TORCH_MINIMUM_VERSION as GLOBAL_MIN_TORCH_VERSION -from .utils import TRANSFORMERS_MINIMUM_VERSION as GLOBAL_MIN_TRANSFORMERS_VERSION -from .utils.doc import add_dynamic_docstring -from .utils.import_utils import is_torch_version, is_transformers_version +from ..utils import TORCH_MINIMUM_VERSION as GLOBAL_MIN_TORCH_VERSION +from ..utils import TRANSFORMERS_MINIMUM_VERSION as GLOBAL_MIN_TRANSFORMERS_VERSION +from ..utils.doc import add_dynamic_docstring +from ..utils.import_utils import is_torch_version, is_transformers_version # from .model_patcher import ModelPatcher From 9eb05a65f3e46fec5cdf5577163210f20e24376c Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Fri, 17 Jan 2025 17:18:23 +0100 Subject: [PATCH 06/22] fix --- optimum/exporters/base.py | 2 +- optimum/exporters/tflite/base.py | 5 +---- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/optimum/exporters/base.py b/optimum/exporters/base.py index 990651c3f4..143728bd3a 100644 --- a/optimum/exporters/base.py +++ b/optimum/exporters/base.py @@ -150,7 +150,7 @@ class ExportersConfig(ABC): def __init__( self, config: "PretrainedConfig", - task: str = "feature-extraction", + task: str, int_dtype: str = "int64", float_dtype: str = "fp32", ): diff --git a/optimum/exporters/tflite/base.py b/optimum/exporters/tflite/base.py index 1c1e725163..907aea53bf 100644 --- a/optimum/exporters/tflite/base.py +++ b/optimum/exporters/tflite/base.py @@ -193,10 +193,7 @@ def __init__( nb_points_per_image: Optional[int] = None, ): super().__init__(config=config, task=task, int_dtype="int64", float_dtype="fp32") - - # self.mandatory_axes = () - # self._axes: Dict[str, int] = {} - + self.task = task # To avoid using **kwargs. axes_values = { "batch_size": batch_size, From 6a34a0c001fee3c53e93a901f84d9d532ba220e0 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Fri, 17 Jan 2025 18:45:26 +0100 Subject: [PATCH 07/22] remove constant --- optimum/exporters/tflite/base.py | 24 +----------------------- 1 file changed, 1 insertion(+), 23 deletions(-) diff --git a/optimum/exporters/tflite/base.py b/optimum/exporters/tflite/base.py index 907aea53bf..25b6a2d098 100644 --- a/optimum/exporters/tflite/base.py +++ b/optimum/exporters/tflite/base.py @@ -148,34 +148,11 @@ class TFLiteConfig(ExportersConfig, ABC): They are required or not depending on the model the `TFLiteConfig` is designed for. """ - NORMALIZED_CONFIG_CLASS: Type = None - DUMMY_INPUT_GENERATOR_CLASSES: Tuple[Type, ...] = () - ATOL_FOR_VALIDATION: Union[float, Dict[str, float]] = 1e-5 MANDATORY_AXES = () SUPPORTED_QUANTIZATION_APPROACHES: Union[ Dict[str, Tuple[QuantizationApproach, ...]], Tuple[QuantizationApproach, ...] ] = tuple(approach for approach in QuantizationApproach) - _TASK_TO_COMMON_OUTPUTS = { - "text-generation": ["logits"], - "feature-extraction": ["last_hidden_state"], - "image-classification": ["logits"], - "image-segmentation": ["logits", "pred_boxes", "pred_masks"], - "masked-im": ["logits"], - "fill-mask": ["logits"], - "multiple-choice": ["logits"], - "object-detection": ["logits", "pred_boxes"], - "question-answering": ["start_logits", "end_logits"], - "semantic-segmentation": ["logits"], - "text2text-generation": ["logits", "encoder_last_hidden_state"], - "text-classification": ["logits"], - "token-classification": ["logits"], - "automatic-speech-recognition": ["logits"], - "audio-classification": ["logits"], - "audio-frame-classification": ["logits"], - "audio-xvector": ["logits"], - } - def __init__( self, config: "PretrainedConfig", @@ -194,6 +171,7 @@ def __init__( ): super().__init__(config=config, task=task, int_dtype="int64", float_dtype="fp32") self.task = task + # To avoid using **kwargs. axes_values = { "batch_size": batch_size, From 655fddf452b255dd0e5633523e048e73b8039c7b Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Fri, 17 Jan 2025 18:46:02 +0100 Subject: [PATCH 08/22] fix style --- optimum/exporters/tflite/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/exporters/tflite/base.py b/optimum/exporters/tflite/base.py index 25b6a2d098..875630aaf7 100644 --- a/optimum/exporters/tflite/base.py +++ b/optimum/exporters/tflite/base.py @@ -19,7 +19,7 @@ from dataclasses import dataclass from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union from transformers.utils import is_tf_available From 6bf67b1b432e39323b4d2c28221cd159eb39a537 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Fri, 17 Jan 2025 19:15:05 +0100 Subject: [PATCH 09/22] move --- optimum/exporters/base.py | 3 --- optimum/exporters/tflite/base.py | 2 ++ 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/optimum/exporters/base.py b/optimum/exporters/base.py index 143728bd3a..ab5041615d 100644 --- a/optimum/exporters/base.py +++ b/optimum/exporters/base.py @@ -160,9 +160,6 @@ def __init__( self.int_dtype = int_dtype self.float_dtype = float_dtype - self.mandatory_axes = () - self._axes: Dict[str, int] = {} - def _create_dummy_input_generator_classes(self, **kwargs) -> List[DummyInputGenerator]: """ Instantiates the dummy input generators from `self.DUMMY_INPUT_GENERATOR_CLASSES`. diff --git a/optimum/exporters/tflite/base.py b/optimum/exporters/tflite/base.py index 875630aaf7..d33a293905 100644 --- a/optimum/exporters/tflite/base.py +++ b/optimum/exporters/tflite/base.py @@ -170,6 +170,8 @@ def __init__( nb_points_per_image: Optional[int] = None, ): super().__init__(config=config, task=task, int_dtype="int64", float_dtype="fp32") + self.mandatory_axes = () + self._axes: Dict[str, int] = {} self.task = task # To avoid using **kwargs. From f97c64d71a212a82a20e138a77fcc09a3a582530 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Fri, 17 Jan 2025 19:20:01 +0100 Subject: [PATCH 10/22] rm comments --- optimum/exporters/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/optimum/exporters/base.py b/optimum/exporters/base.py index ab5041615d..91fa78487e 100644 --- a/optimum/exporters/base.py +++ b/optimum/exporters/base.py @@ -128,7 +128,7 @@ class ExportersConfig(ABC): "feature-extraction": ["last_hidden_state"], # for neuron : ["last_hidden_state", "pooler_output"] "fill-mask": ["logits"], "image-classification": ["logits"], - "image-segmentation": ["logits"], # for tflite : ["logits", "pred_boxes", "pred_masks"] + "image-segmentation": ["logits"], "image-to-text": ["logits"], "image-to-image": ["reconstruction"], "mask-generation": ["logits"], @@ -137,7 +137,7 @@ class ExportersConfig(ABC): "object-detection": ["logits", "pred_boxes"], "question-answering": ["start_logits", "end_logits"], "semantic-segmentation": ["logits"], - "text2text-generation": ["logits"], # for tflite : ["logits", "encoder_last_hidden_state"], + "text2text-generation": ["logits"], "text-classification": ["logits"], "text-generation": ["logits"], "time-series-forecasting": ["prediction_outputs"], From cfef7b4db6ee294211b20d70c8c937be591ff5a6 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Tue, 28 Jan 2025 17:48:12 +0100 Subject: [PATCH 11/22] fix masked-im output names --- optimum/exporters/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/exporters/base.py b/optimum/exporters/base.py index 91fa78487e..58a89d5992 100644 --- a/optimum/exporters/base.py +++ b/optimum/exporters/base.py @@ -132,7 +132,7 @@ class ExportersConfig(ABC): "image-to-text": ["logits"], "image-to-image": ["reconstruction"], "mask-generation": ["logits"], - "masked-im": ["logits"], # for onnx : ["reconstruction"] + "masked-im": ["reconstruction"], "multiple-choice": ["logits"], "object-detection": ["logits", "pred_boxes"], "question-answering": ["start_logits", "end_logits"], From 14aec3f93ba664b6efc2dca8c6052184a7f8d1f6 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Tue, 28 Jan 2025 18:23:57 +0100 Subject: [PATCH 12/22] style --- optimum/exporters/base.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/optimum/exporters/base.py b/optimum/exporters/base.py index 58a89d5992..c751d8ea2c 100644 --- a/optimum/exporters/base.py +++ b/optimum/exporters/base.py @@ -20,10 +20,6 @@ from transformers.utils import is_torch_available - -if is_torch_available(): - pass - from ..utils import ( DEFAULT_DUMMY_SHAPES, DummyInputGenerator, @@ -167,7 +163,6 @@ def _create_dummy_input_generator_classes(self, **kwargs) -> List[DummyInputGene forces the other generators to use the same batch size, meaning they will all produce inputs of the same batch size. Override this method for custom behavior. """ - # self._validate_mandatory_axes() return [cls_(self.task, self._normalized_config, **kwargs) for cls_ in self.DUMMY_INPUT_GENERATOR_CLASSES] @property From 7d46f0174b58d7d0aa2c7959c317bd393a00ca51 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Tue, 28 Jan 2025 18:24:45 +0100 Subject: [PATCH 13/22] rm comment --- optimum/exporters/base.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/optimum/exporters/base.py b/optimum/exporters/base.py index c751d8ea2c..d2e7eb7b73 100644 --- a/optimum/exporters/base.py +++ b/optimum/exporters/base.py @@ -31,8 +31,6 @@ from ..utils.import_utils import is_torch_version, is_transformers_version -# from .model_patcher import ModelPatcher - if TYPE_CHECKING: from transformers import PretrainedConfig @@ -265,12 +263,3 @@ def flatten_inputs(cls, inputs: Dict[str, Any]) -> Dict[str, Any]: else: flatten[name] = value return flatten - - ############################################################################################################################################################ - - # def patch_model_for_export( - # self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - # ) -> ModelPatcher: - # return ModelPatcher(self, model, model_kwargs=model_kwargs) - - ############################################################################################################################################################ From f98da02fbc5ecf61dd1b532cd3372f49dec44031 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Wed, 29 Jan 2025 16:49:07 +0100 Subject: [PATCH 14/22] fix doc generation --- optimum/exporters/base.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/optimum/exporters/base.py b/optimum/exporters/base.py index d2e7eb7b73..6ddd4c5065 100644 --- a/optimum/exporters/base.py +++ b/optimum/exporters/base.py @@ -66,7 +66,7 @@ The number of frames to use in the dummpy inputs for audio tasks in case the input is raw audio. Returns: - `Dict`: A dictionary mapping the input names to dummy tensors in the proper framework format. + `Dict[str, [tf.Tensor, torch.Tensor]]`: A dictionary mapping the input names to dummy tensors in the proper framework format. """ @@ -224,14 +224,6 @@ def is_torch_support_available(self) -> bool: @add_dynamic_docstring(text=GENERATE_DUMMY_DOCSTRING, dynamic_elements=DEFAULT_DUMMY_SHAPES) def generate_dummy_inputs(self, framework: str = "pt", **kwargs) -> Dict: - """ - Generates dummy inputs that the exported model should be able to process. - This method is actually used to determine the input specs that are needed for the export. - - Returns: - `Dict[str, [tf.Tensor, torch.Tensor]]`: A dictionary mapping input names to dummy tensors. - """ - dummy_inputs_generators = self._create_dummy_input_generator_classes(**kwargs) dummy_inputs = {} for input_name in self.inputs: From bc423a5b8f1d537c4bfa3473ba4652947604d027 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Wed, 29 Jan 2025 17:50:54 +0100 Subject: [PATCH 15/22] task attribute already set --- optimum/exporters/tflite/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/optimum/exporters/tflite/base.py b/optimum/exporters/tflite/base.py index d33a293905..2079ff372e 100644 --- a/optimum/exporters/tflite/base.py +++ b/optimum/exporters/tflite/base.py @@ -169,10 +169,10 @@ def __init__( point_batch_size: Optional[int] = None, nb_points_per_image: Optional[int] = None, ): - super().__init__(config=config, task=task, int_dtype="int64", float_dtype="fp32") self.mandatory_axes = () self._axes: Dict[str, int] = {} - self.task = task + + super().__init__(config=config, task=task, int_dtype="int64", float_dtype="fp32") # To avoid using **kwargs. axes_values = { From bb259cd29d4a701f1142445e1946ed3cd2d08c3f Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Fri, 31 Jan 2025 17:59:59 +0100 Subject: [PATCH 16/22] model patcher --- optimum/exporters/onnx/base.py | 8 +++++++- optimum/exporters/onnx/config.py | 8 -------- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py index 00e126c2f2..66da6156c1 100644 --- a/optimum/exporters/onnx/base.py +++ b/optimum/exporters/onnx/base.py @@ -49,7 +49,7 @@ ) from ..base import ExportersConfig from .constants import ONNX_DECODER_MERGED_NAME, ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME -from .model_patcher import ModelPatcher, Seq2SeqModelPatcher +from .model_patcher import DecoderModelPatcher, ModelPatcher, Seq2SeqModelPatcher # TODO : moved back onnx imports applied in https://github.com/huggingface/optimum/pull/2114/files after refactorization @@ -627,6 +627,12 @@ def generate_dummy_inputs_for_validation( return reference_model_inputs + def patch_model_for_export( + self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None + ) -> "ModelPatcher": + # Refer to DecoderModelPatcher. + return DecoderModelPatcher(self, model, model_kwargs=model_kwargs) + class ConfigBehavior(str, enum.Enum): """ diff --git a/optimum/exporters/onnx/config.py b/optimum/exporters/onnx/config.py index 69366d6be1..6b50f3162d 100644 --- a/optimum/exporters/onnx/config.py +++ b/optimum/exporters/onnx/config.py @@ -34,7 +34,6 @@ ) from .base import ConfigBehavior, OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast from .constants import ONNX_DECODER_MERGED_NAME, ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME -from .model_patcher import DecoderModelPatcher # TODO : moved back onnx imports applied in https://github.com/huggingface/optimum/pull/2114/files after refactorization @@ -43,7 +42,6 @@ if TYPE_CHECKING: from transformers import PretrainedConfig, PreTrainedModel - from .model_patcher import ModelPatcher if is_tf_available(): from transformers import TFPreTrainedModel @@ -160,12 +158,6 @@ def post_process_exported_models( return models_and_onnx_configs, onnx_files_subpaths - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> "ModelPatcher": - # Refer to DecoderModelPatcher. - return DecoderModelPatcher(self, model, model_kwargs=model_kwargs) - class TextDecoderWithPositionIdsOnnxConfig(TextDecoderOnnxConfig): @property From 1fbaf644f61504d8b6431e796e62092f32f5f53a Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Fri, 31 Jan 2025 18:21:31 +0100 Subject: [PATCH 17/22] add model patcher --- optimum/exporters/base.py | 9 +- optimum/exporters/onnx/base.py | 20 +---- optimum/exporters/onnx/config.py | 1 - optimum/exporters/onnx/model_configs.py | 107 +++++------------------- 4 files changed, 34 insertions(+), 103 deletions(-) diff --git a/optimum/exporters/base.py b/optimum/exporters/base.py index 6ddd4c5065..2d74e90324 100644 --- a/optimum/exporters/base.py +++ b/optimum/exporters/base.py @@ -29,10 +29,11 @@ from ..utils import TRANSFORMERS_MINIMUM_VERSION as GLOBAL_MIN_TRANSFORMERS_VERSION from ..utils.doc import add_dynamic_docstring from ..utils.import_utils import is_torch_version, is_transformers_version +from .model_patcher import ModelPatcher if TYPE_CHECKING: - from transformers import PretrainedConfig + from transformers import PretrainedConfig, PreTrainedModel, TFPreTrainedModel from .model_patcher import PatchingSpec @@ -140,6 +141,7 @@ class ExportersConfig(ABC): "zero-shot-image-classification": ["logits_per_image", "logits_per_text", "text_embeds", "image_embeds"], "zero-shot-object-detection": ["logits", "pred_boxes", "text_embeds", "image_embeds"], } + _MODEL_PATCHER = ModelPatcher def __init__( self, @@ -255,3 +257,8 @@ def flatten_inputs(cls, inputs: Dict[str, Any]) -> Dict[str, Any]: else: flatten[name] = value return flatten + + def patch_model_for_export( + self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None + ) -> "ModelPatcher": + return self._MODEL_PATCHER(self, model, model_kwargs=model_kwargs) diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py index 66da6156c1..4246cb2a21 100644 --- a/optimum/exporters/onnx/base.py +++ b/optimum/exporters/onnx/base.py @@ -49,7 +49,7 @@ ) from ..base import ExportersConfig from .constants import ONNX_DECODER_MERGED_NAME, ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME -from .model_patcher import DecoderModelPatcher, ModelPatcher, Seq2SeqModelPatcher +from .model_patcher import DecoderModelPatcher, Seq2SeqModelPatcher # TODO : moved back onnx imports applied in https://github.com/huggingface/optimum/pull/2114/files after refactorization @@ -272,11 +272,6 @@ def fix_dynamic_axes( del onnx_model gc.collect() - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> ModelPatcher: - return ModelPatcher(self, model, model_kwargs=model_kwargs) - @property def torch_to_onnx_input_map(self) -> Dict[str, str]: """ @@ -435,6 +430,7 @@ class OnnxConfigWithPast(OnnxConfig, ABC): PAD_ATTENTION_MASK_TO_PAST: bool = False SUPPORTS_PAST: bool = True + _MODEL_PATCHER = DecoderModelPatcher def __init__( self, @@ -627,12 +623,6 @@ def generate_dummy_inputs_for_validation( return reference_model_inputs - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> "ModelPatcher": - # Refer to DecoderModelPatcher. - return DecoderModelPatcher(self, model, model_kwargs=model_kwargs) - class ConfigBehavior(str, enum.Enum): """ @@ -653,6 +643,7 @@ class OnnxSeq2SeqConfigWithPast(OnnxConfigWithPast): """ DUMMY_PKV_GENERATOR_CLASS = DummySeq2SeqPastKeyValuesGenerator + _MODEL_PATCHER = Seq2SeqModelPatcher def __init__( self, @@ -785,11 +776,6 @@ def flatten_past_key_values(self, flattened_output, name, idx, t): flattened_output[f"{name}.{idx}.encoder.key"] = t[2] flattened_output[f"{name}.{idx}.encoder.value"] = t[3] - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> ModelPatcher: - return Seq2SeqModelPatcher(self, model, model_kwargs=model_kwargs) - def post_process_exported_models( self, path: Path, diff --git a/optimum/exporters/onnx/config.py b/optimum/exporters/onnx/config.py index 6b50f3162d..290e05d2b3 100644 --- a/optimum/exporters/onnx/config.py +++ b/optimum/exporters/onnx/config.py @@ -42,7 +42,6 @@ if TYPE_CHECKING: from transformers import PretrainedConfig, PreTrainedModel - if is_tf_available(): from transformers import TFPreTrainedModel diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py index 78e1ca418d..755d1b58d4 100644 --- a/optimum/exporters/onnx/model_configs.py +++ b/optimum/exporters/onnx/model_configs.py @@ -102,7 +102,6 @@ from transformers import PretrainedConfig from transformers.modeling_utils import PreTrainedModel - from .model_patcher import ModelPatcher if is_tf_available(): from transformers.modeling_tf_utils import TFPreTrainedModel @@ -395,11 +394,7 @@ class MistralOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): ) + TextDecoderOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_key_value_heads="num_key_value_heads", allow_new=True) - - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> "ModelPatcher": - return MistralModelPatcher(self, model, model_kwargs=model_kwargs) + _MODEL_PATCHER = MistralModelPatcher class MPTOnnxConfig(TextDecoderOnnxConfig): @@ -487,6 +482,10 @@ class FalconOnnxConfig(TextDecoderOnnxConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig DUMMY_PKV_GENERATOR_CLASS = FalconDummyPastKeyValuesGenerator + # we need to set output_attentions=True in the model input to avoid calling + # torch.nn.functional.scaled_dot_product_attention that is not supported by the ONNX export + _MODEL_PATCHER = FalconModelPatcher + def __init__( self, config: "PretrainedConfig", @@ -527,13 +526,6 @@ def inputs(self) -> Dict[str, Dict[int, str]]: return common_inputs - # we need to set output_attentions=True in the model input to avoid calling - # torch.nn.functional.scaled_dot_product_attention that is not supported by the ONNX export - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> "ModelPatcher": - return FalconModelPatcher(self, model, model_kwargs=model_kwargs) - class T5DummySeq2SeqPastKeyValuesGenerator(DummySeq2SeqPastKeyValuesGenerator): def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): @@ -1040,6 +1032,8 @@ def torch_to_onnx_input_map(self) -> Dict[str, str]: class MgpstrOnnxConfig(ViTOnnxConfig): + _MODEL_PATCHER = MgpstrModelPatcher + @property def outputs(self) -> Dict[str, Dict[int, str]]: return { @@ -1048,15 +1042,14 @@ def outputs(self) -> Dict[str, Dict[int, str]]: "wp_logits": {0: "batch_size"}, } - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> "ModelPatcher": - return MgpstrModelPatcher(self, model, model_kwargs=model_kwargs) - class SentenceTransformersTransformerOnnxConfig(TextEncoderOnnxConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig DEFAULT_ONNX_OPSET = 14 # Some bottleneck transformers models require a specific ONNX opset to be successfully exported. We put a rather high opset here for the export to work for all architectures. + # we need to set output_attentions=True in the model input to avoid calling + # torch.nn.functional.scaled_dot_product_attention that is not supported by the ONNX export + # due to the op torch.nn.functional.multi_head_attention_forward used for WavLM + _MODEL_PATCHER = SentenceTransformersTransformerPatcher @property def inputs(self) -> Dict[str, Dict[int, str]]: @@ -1072,14 +1065,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: "sentence_embedding": {0: "batch_size"}, } - # we need to set output_attentions=True in the model input to avoid calling - # torch.nn.functional.scaled_dot_product_attention that is not supported by the ONNX export - # due to the op torch.nn.functional.multi_head_attention_forward used for WavLM - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> "ModelPatcher": - return SentenceTransformersTransformerPatcher(self, model, model_kwargs=model_kwargs) - class CLIPNormalizedConfig(NormalizedTextAndVisionConfig): TEXT_CONFIG = "text_config" @@ -1088,6 +1073,7 @@ class CLIPNormalizedConfig(NormalizedTextAndVisionConfig): class CLIPVisionModelOnnxConfig(VisionOnnxConfig): NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig + _MODEL_PATCHER = CLIPModelPatcher @property def inputs(self) -> Dict[str, Dict[int, str]]: @@ -1101,16 +1087,10 @@ def outputs(self) -> Dict[str, Dict[int, str]]: return common_outputs - def patch_model_for_export( - self, - model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], - model_kwargs: Optional[Dict[str, Any]] = None, - ) -> "ModelPatcher": - return CLIPModelPatcher(self, model, model_kwargs=model_kwargs) - class CLIPOnnxConfig(TextAndVisionOnnxConfig): NORMALIZED_CONFIG_CLASS = CLIPNormalizedConfig + _MODEL_PATCHER = CLIPModelPatcher @property def inputs(self) -> Dict[str, Dict[int, str]]: @@ -1129,15 +1109,10 @@ def outputs(self) -> Dict[str, Dict[int, str]]: "image_embeds": {0: "image_batch_size"}, } - def patch_model_for_export( - self, - model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], - model_kwargs: Optional[Dict[str, Any]] = None, - ) -> "ModelPatcher": - return CLIPModelPatcher(self, model, model_kwargs=model_kwargs) - class SentenceTransformersCLIPOnnxConfig(CLIPOnnxConfig): + _MODEL_PATCHER = SentenceTransformersCLIPPatcher + @property def outputs(self) -> Dict[str, Dict[int, str]]: return { @@ -1145,11 +1120,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: "image_embeds": {0: "image_batch_size"}, } - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> "ModelPatcher": - return SentenceTransformersCLIPPatcher(self, model, model_kwargs=model_kwargs) - class CLIPTextWithProjectionOnnxConfig(TextEncoderOnnxConfig): ATOL_FOR_VALIDATION = 1e-3 @@ -1162,6 +1132,7 @@ class CLIPTextWithProjectionOnnxConfig(TextEncoderOnnxConfig): num_layers="num_hidden_layers", allow_new=True, ) + _MODEL_PATCHER = CLIPModelPatcher @property def inputs(self) -> Dict[str, Dict[int, str]]: @@ -1181,15 +1152,10 @@ def outputs(self) -> Dict[str, Dict[int, str]]: return common_outputs - def patch_model_for_export( - self, - model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], - model_kwargs: Optional[Dict[str, Any]] = None, - ) -> "ModelPatcher": - return CLIPModelPatcher(self, model, model_kwargs=model_kwargs) - class CLIPTextOnnxConfig(CLIPTextWithProjectionOnnxConfig): + _MODEL_PATCHER = CLIPModelPatcher + @property def outputs(self) -> Dict[str, Dict[int, str]]: common_outputs = { @@ -1203,13 +1169,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: return common_outputs - def patch_model_for_export( - self, - model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], - model_kwargs: Optional[Dict[str, Any]] = None, - ) -> "ModelPatcher": - return CLIPModelPatcher(self, model, model_kwargs=model_kwargs) - class SiglipNormalizedConfig(CLIPNormalizedConfig): pass @@ -1707,14 +1666,10 @@ class UniSpeechSATOnnxConfig(HubertOnnxConfig): class WavLMOnnxConfig(HubertOnnxConfig): DEFAULT_ONNX_OPSET = 12 - # we need to set output_attentions=True in the model input to avoid calling # torch.nn.functional.scaled_dot_product_attention that is not supported by the ONNX export # due to the op torch.nn.functional.multi_head_attention_forward used for WavLM - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> "ModelPatcher": - return WavLMModelPatcher(self, model, model_kwargs=model_kwargs) + _MODEL_PATCHER = WavLMModelPatcher class ASTDummyAudioInputGenerator(DummyAudioInputGenerator): @@ -1823,6 +1778,7 @@ class MusicgenOnnxConfig(OnnxSeq2SeqConfigWithPast): DummyIntGenerator, ) DUMMY_PKV_GENERATOR_CLASS = DummySeq2SeqPastKeyValuesGenerator + _MODEL_PATCHER = MusicgenModelPatcher def __init__( self, @@ -1999,11 +1955,6 @@ def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], dire 2: "encoder_sequence_length_out", } - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> "ModelPatcher": - return MusicgenModelPatcher(self, model, model_kwargs=model_kwargs) - @property def torch_to_onnx_input_map(self) -> Dict[str, str]: if self._behavior is ConfigBehavior.DECODER: @@ -2115,6 +2066,7 @@ class SpeechT5OnnxConfig(OnnxSeq2SeqConfigWithPast): "without-past": "The same as `with-past`, just without KV cache support. This is not a recommended export as slower than `with-past`.", } DEFAULT_VARIANT = "with-past" + _MODEL_PATCHER = SpeechT5ModelPatcher def __init__( self, @@ -2195,11 +2147,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: return common_outputs - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> "ModelPatcher": - return SpeechT5ModelPatcher(self, model, model_kwargs=model_kwargs) - @property def torch_to_onnx_input_map(self) -> Dict[str, str]: return {"encoder_outputs": "encoder_hidden_states"} @@ -2333,6 +2280,7 @@ class VisionEncoderDecoderOnnxConfig(EncoderDecoderBaseOnnxConfig): DEFAULT_ONNX_OPSET = 14 # uses SDPA in Transformers, hence opset>=14. DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyVisionEncoderDecoderPastKeyValuesGenerator) + _MODEL_PATCHER = VisionEncoderDecoderPatcher @property def inputs(self) -> Dict[str, Dict[int, str]]: @@ -2365,11 +2313,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: # so we can not initializer MBartONNXConfig with document-question-answering). return super().outputs - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> "ModelPatcher": - return VisionEncoderDecoderPatcher(self, model, model_kwargs=model_kwargs) - class SamOnnxConfig(OnnxConfig): MIN_TRANSFORMERS_VERSION = version.parse("4.29.0.dev0") @@ -2383,6 +2326,7 @@ class SamOnnxConfig(OnnxConfig): "split": "The vision encoder is exported as a separate vision_encoder.onnx, and the prompt encoder and mask decoder are exported as a prompt_encoder_mask_decoder.onnx. This allows to encoder the image only once for multiple point queries.", } DEFAULT_VARIANT = "split" + _MODEL_PATCHER = SAMModelPatcher def __init__( self, @@ -2437,11 +2381,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: "pred_masks": {0: "batch_size", 1: "point_batch_size"}, } - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> "ModelPatcher": - return SAMModelPatcher(self, model, model_kwargs=model_kwargs) - class Pix2StructNormalizedConfig(NormalizedSeq2SeqConfig): ENCODER_NUM_LAYERS = "vision_config.num_hidden_layers" From 795a618b5e4b14fe1dae7e101c892fc5b25faa52 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Wed, 12 Feb 2025 17:55:06 +0100 Subject: [PATCH 18/22] add processor --- optimum/exporters/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/exporters/utils.py b/optimum/exporters/utils.py index aae07c411d..58e170ba97 100644 --- a/optimum/exporters/utils.py +++ b/optimum/exporters/utils.py @@ -545,7 +545,7 @@ def get_speecht5_models_for_export( use_past=use_past, use_past_in_inputs=False, # Irrelevant here. behavior=config._behavior, # Irrelevant here. - # preprocessors=config._preprocessors, + preprocessors=config._preprocessors, is_postnet_and_vocoder=True, legacy=config.legacy, ) From 7a3f3389d0c7a331900d111686ab46849afd0d3b Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Wed, 12 Feb 2025 17:59:40 +0100 Subject: [PATCH 19/22] move patching specs to onnx config --- optimum/exporters/base.py | 3 --- optimum/exporters/onnx/base.py | 5 ++++- optimum/exporters/onnx/model_configs.py | 1 - 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/optimum/exporters/base.py b/optimum/exporters/base.py index 2d74e90324..08fed1e47e 100644 --- a/optimum/exporters/base.py +++ b/optimum/exporters/base.py @@ -35,8 +35,6 @@ if TYPE_CHECKING: from transformers import PretrainedConfig, PreTrainedModel, TFPreTrainedModel - from .model_patcher import PatchingSpec - logger = logging.get_logger(__name__) @@ -112,7 +110,6 @@ class ExportersConfig(ABC): ATOL_FOR_VALIDATION: Union[float, Dict[str, float]] = 1e-5 MIN_TORCH_VERSION = GLOBAL_MIN_TORCH_VERSION MIN_TRANSFORMERS_VERSION = GLOBAL_MIN_TRANSFORMERS_VERSION - PATCHING_SPECS: Optional[List["PatchingSpec"]] = None _TASK_TO_COMMON_OUTPUTS = { "audio-classification": ["logits"], "audio-frame-classification": ["logits"], diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py index 4336665dcb..96659b786d 100644 --- a/optimum/exporters/onnx/base.py +++ b/optimum/exporters/onnx/base.py @@ -60,6 +60,8 @@ if TYPE_CHECKING: from transformers import PretrainedConfig, PreTrainedModel, TFPreTrainedModel + from .model_patcher import PatchingSpec + if is_diffusers_available(): from diffusers import ModelMixin @@ -102,7 +104,8 @@ class OnnxConfig(ExportersConfig): DEFAULT_ONNX_OPSET = 11 VARIANTS = {"default": "The default ONNX variant."} DEFAULT_VARIANT = "default" - + # TODO: move PATCHING_SPECS to ExportersConfig + PATCHING_SPECS: Optional[List["PatchingSpec"]] = None _TASK_TO_COMMON_OUTPUTS = { "audio-classification": OrderedDict({"logits": {0: "batch_size"}}), "audio-frame-classification": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py index 64499352dc..0c51de3d35 100644 --- a/optimum/exporters/onnx/model_configs.py +++ b/optimum/exporters/onnx/model_configs.py @@ -104,7 +104,6 @@ from transformers import PretrainedConfig from transformers.modeling_utils import PreTrainedModel - if is_tf_available(): from transformers.modeling_tf_utils import TFPreTrainedModel From 65ad611d92ca24da663a03c3f9199da1b0020163 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Wed, 12 Feb 2025 18:07:01 +0100 Subject: [PATCH 20/22] rename --- optimum/exporters/base.py | 5 ----- optimum/exporters/onnx/base.py | 6 +++--- optimum/exporters/tflite/base.py | 4 ++-- 3 files changed, 5 insertions(+), 10 deletions(-) diff --git a/optimum/exporters/base.py b/optimum/exporters/base.py index 08fed1e47e..db84d0fa27 100644 --- a/optimum/exporters/base.py +++ b/optimum/exporters/base.py @@ -69,12 +69,7 @@ """ -# TODO: Remove class ExportConfig(ABC): - pass - - -class ExportersConfig(ABC): """ Base class describing metadata on how to export the model through the ONNX format. diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py index 96659b786d..d8f5b587a6 100644 --- a/optimum/exporters/onnx/base.py +++ b/optimum/exporters/onnx/base.py @@ -47,7 +47,7 @@ is_onnxruntime_available, is_transformers_version, ) -from ..base import ExportersConfig +from ..base import ExportConfig from .constants import ONNX_DECODER_MERGED_NAME, ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME from .model_patcher import DecoderModelPatcher, Seq2SeqModelPatcher @@ -100,11 +100,11 @@ """ -class OnnxConfig(ExportersConfig): +class OnnxConfig(ExportConfig): DEFAULT_ONNX_OPSET = 11 VARIANTS = {"default": "The default ONNX variant."} DEFAULT_VARIANT = "default" - # TODO: move PATCHING_SPECS to ExportersConfig + # TODO: move PATCHING_SPECS to ExportConfig PATCHING_SPECS: Optional[List["PatchingSpec"]] = None _TASK_TO_COMMON_OUTPUTS = { "audio-classification": OrderedDict({"logits": {0: "batch_size"}}), diff --git a/optimum/exporters/tflite/base.py b/optimum/exporters/tflite/base.py index 2079ff372e..5780751cb0 100644 --- a/optimum/exporters/tflite/base.py +++ b/optimum/exporters/tflite/base.py @@ -27,7 +27,7 @@ if is_tf_available(): import tensorflow as tf -from ..base import ExportersConfig +from ..base import ExportConfig if TYPE_CHECKING: @@ -115,7 +115,7 @@ def __post_init__(self): self.approach = QuantizationApproach(self.approach) -class TFLiteConfig(ExportersConfig, ABC): +class TFLiteConfig(ExportConfig, ABC): """ Base class for TFLite exportable model describing metadata on how to export the model through the TFLite format. From 46aac954c4db7a327e37588d30aaeae07027c1ac Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Wed, 12 Feb 2025 18:12:29 +0100 Subject: [PATCH 21/22] add comment --- optimum/exporters/base.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/optimum/exporters/base.py b/optimum/exporters/base.py index db84d0fa27..63246f9387 100644 --- a/optimum/exporters/base.py +++ b/optimum/exporters/base.py @@ -29,11 +29,10 @@ from ..utils import TRANSFORMERS_MINIMUM_VERSION as GLOBAL_MIN_TRANSFORMERS_VERSION from ..utils.doc import add_dynamic_docstring from ..utils.import_utils import is_torch_version, is_transformers_version -from .model_patcher import ModelPatcher if TYPE_CHECKING: - from transformers import PretrainedConfig, PreTrainedModel, TFPreTrainedModel + from transformers import PretrainedConfig logger = logging.get_logger(__name__) @@ -133,7 +132,8 @@ class ExportConfig(ABC): "zero-shot-image-classification": ["logits_per_image", "logits_per_text", "text_embeds", "image_embeds"], "zero-shot-object-detection": ["logits", "pred_boxes", "text_embeds", "image_embeds"], } - _MODEL_PATCHER = ModelPatcher + # TODO : add _MODEL_PATCHER + patch_model_for_export + # _MODEL_PATCHER = ModelPatcher def __init__( self, @@ -249,8 +249,3 @@ def flatten_inputs(cls, inputs: Dict[str, Any]) -> Dict[str, Any]: else: flatten[name] = value return flatten - - def patch_model_for_export( - self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None - ) -> "ModelPatcher": - return self._MODEL_PATCHER(self, model, model_kwargs=model_kwargs) From e7d8c1e76e9637435c3ac08b0f3c88b3401b4084 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Fri, 7 Mar 2025 18:10:51 +0100 Subject: [PATCH 22/22] fix --- optimum/exporters/onnx/base.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py index d8f5b587a6..b64de5df28 100644 --- a/optimum/exporters/onnx/base.py +++ b/optimum/exporters/onnx/base.py @@ -49,7 +49,7 @@ ) from ..base import ExportConfig from .constants import ONNX_DECODER_MERGED_NAME, ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME -from .model_patcher import DecoderModelPatcher, Seq2SeqModelPatcher +from .model_patcher import DecoderModelPatcher, ModelPatcher, Seq2SeqModelPatcher # TODO : moved back onnx imports applied in https://github.com/huggingface/optimum/pull/2114/files after refactorization @@ -100,12 +100,14 @@ """ -class OnnxConfig(ExportConfig): +class OnnxConfig(ExportConfig, ABC): DEFAULT_ONNX_OPSET = 11 VARIANTS = {"default": "The default ONNX variant."} DEFAULT_VARIANT = "default" # TODO: move PATCHING_SPECS to ExportConfig PATCHING_SPECS: Optional[List["PatchingSpec"]] = None + _MODEL_PATCHER = ModelPatcher + _TASK_TO_COMMON_OUTPUTS = { "audio-classification": OrderedDict({"logits": {0: "batch_size"}}), "audio-frame-classification": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), @@ -425,6 +427,11 @@ def post_process_exported_models( return models_and_onnx_configs, onnx_files_subpaths + def patch_model_for_export( + self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None + ) -> ModelPatcher: + return self._MODEL_PATCHER(self, model, model_kwargs=model_kwargs) + class OnnxConfigWithPast(OnnxConfig, ABC): """