huggingface · fxmarty · Feb 15, 2023 · Dec 27, 2022 · Dec 27, 2022 · Dec 27, 2022
diff --git a/.github/workflows/test_export_onnx_cli.yml b/.github/workflows/test_export_onnx_cli.yml
@@ -27,7 +27,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        pip install .[tests,exporters-tf]
+        pip install .[tests,exporters]
     - name: Test with unittest
       working-directory: tests
       run: |

diff --git a/optimum/commands/export/onnx.py b/optimum/commands/export/onnx.py
@@ -90,7 +90,15 @@ def parse_args_onnx(parser):
     optional_group.add_argument(
         "--trust-remote-code",
         action="store_true",
-        help="Allow to use custom code for the modeling hosted in the model repository. This option should only be set for repositories you trust and in which you have read the code, as it will execute on your local machine arbitrary code present in the model repository.",
+        help="Allows to use custom code for the modeling hosted in the model repository. This option should only be set for repositories you trust and in which you have read the code, as it will execute on your local machine arbitrary code present in the model repository.",
+    )
+    optional_group.add_argument(
+        "--no-post-process",
+        action="store_true",
+        help=(
+            "Allows to disable any post-processing done by default on the exported ONNX models. For example, the merging of decoder"
+            " and decoder-with-past models into a single ONNX model file to reduce memory usage."
+        ),
     )
 
     input_group = parser.add_argument_group(

diff --git a/optimum/exporters/onnx/__main__.py b/optimum/exporters/onnx/__main__.py
@@ -24,7 +24,7 @@
 from ..error_utils import AtolError, OutputMatchError, ShapeError
 from ..tasks import TasksManager
 from .base import OnnxConfigWithPast
-from .convert import export, export_models, validate_model_outputs, validate_models_outputs
+from .convert import export_models, validate_models_outputs
 from .utils import (
     get_decoder_models_for_export,
     get_encoder_decoder_models_for_export,
@@ -43,10 +43,9 @@ def main():
 
     # Retrieve CLI arguments
     args = parser.parse_args()
-    args.output = args.output.joinpath("model.onnx")
 
-    if not args.output.parent.exists():
-        args.output.parent.mkdir(parents=True)
+    if not args.output.exists():
+        args.output.mkdir(parents=True)
 
     if args.for_ort:
         logger.warning(
@@ -88,10 +87,10 @@ def main():
         else:
             logger.info(
                 f"The task `{task}` was manually specified, and past key values will not be reused in the decoding."
-                f"Please pass `--task {task}-with-past` to export using the past key values."
+                f" if needed, please pass `--task {task}-with-past` to export using the past key values."
             )
 
-    if task == "auto":
+    if args.task == "auto":
         logger.info(f"Automatic task detection to {task}.")
 
     if task != "stable-diffusion":
@@ -130,96 +129,87 @@ def main():
                 args.atol = args.atol[task.replace("-with-past", "")]
 
         # Saving the model config and preprocessor as this is needed sometimes.
-        model.config.save_pretrained(args.output.parent)
-        maybe_save_preprocessors(args.model, args.output.parent)
-
-    if task == "stable-diffusion" or (
-        task.startswith(("causal-lm", "seq2seq-lm", "speech2seq-lm", "vision2seq-lm", "default-with-past"))
-        and not args.monolith
-    ):
-        if task == "stable-diffusion":
-            output_names = [
-                "text_encoder/model.onnx",
-                "unet/model.onnx",
-                "vae_encoder/model.onnx",
-                "vae_decoder/model.onnx",
-            ]
-            models_and_onnx_configs = get_stable_diffusion_models_for_export(model)
-            # Saving the additional components needed to perform inference.
-            model.tokenizer.save_pretrained(args.output.parent.joinpath("tokenizer"))
-            model.scheduler.save_pretrained(args.output.parent.joinpath("scheduler"))
-            model.feature_extractor.save_pretrained(args.output.parent.joinpath("feature_extractor"))
-            model.save_config(args.output.parent)
-        else:
-            if model.config.is_encoder_decoder and task.startswith("causal-lm"):
-                raise ValueError(
-                    f"model.config.is_encoder_decoder is True and task is `{task}`, which are incompatible. If the task was auto-inferred, please fill a bug report"
-                    f"at https://github.com/huggingface/optimum, if --task was explicitely passed, make sure you selected the right task for the model,"
-                    f" referring to `optimum.exporters.tasks.TaskManager`'s `_TASKS_TO_AUTOMODELS`."
-                )
-            if model.config.is_encoder_decoder:
-                models_and_onnx_configs = get_encoder_decoder_models_for_export(model, onnx_config)
-            else:
-                models_and_onnx_configs = get_decoder_models_for_export(model, onnx_config)
-            output_names = None
-
-        onnx_inputs, onnx_outputs = export_models(
-            models_and_onnx_configs=models_and_onnx_configs,
-            opset=args.opset,
-            output_dir=args.output.parent,
-            output_names=output_names,
-            input_shapes=input_shapes,
-            device=args.device,
-        )
+        model.config.save_pretrained(args.output)
+        maybe_save_preprocessors(args.model, args.output)
+
+    if task == "stable-diffusion":
+        onnx_files_subpaths = [
+            "text_encoder/model.onnx",
+            "unet/model.onnx",
+            "vae_encoder/model.onnx",
+            "vae_decoder/model.onnx",
+        ]
+        models_and_onnx_configs = get_stable_diffusion_models_for_export(model)
+        # Saving the additional components needed to perform inference.
+        model.tokenizer.save_pretrained(args.output.joinpath("tokenizer"))
+        model.scheduler.save_pretrained(args.output.joinpath("scheduler"))
+        model.feature_extractor.save_pretrained(args.output.joinpath("feature_extractor"))
+        model.save_config(args.output)
     else:
-        onnx_inputs, onnx_outputs = export(
-            model=model,
-            config=onnx_config,
-            output=args.output,
-            opset=args.opset,
-            input_shapes=input_shapes,
-            device=args.device,
-        )
+        if model.config.is_encoder_decoder and task.startswith("causal-lm"):
+            raise ValueError(
+                f"model.config.is_encoder_decoder is True and task is `{task}`, which are incompatible. If the task was auto-inferred, please fill a bug report"
+                f"at https://github.com/huggingface/optimum, if --task was explicitely passed, make sure you selected the right task for the model,"
+                f" referring to `optimum.exporters.tasks.TaskManager`'s `_TASKS_TO_AUTOMODELS`."
+            )
 
-    try:
-        if task == "stable-diffusion" or (
-            task.startswith(("causal-lm", "seq2seq-lm", "speech2seq-lm", "vision2seq-lm", "default-with-past"))
+        onnx_files_subpaths = None
+        if (
+            model.config.is_encoder_decoder
+            and task.startswith(("seq2seq-lm", "speech2seq-lm", "vision2seq-lm", "default-with-past"))
             and not args.monolith
         ):
-            validate_models_outputs(
-                models_and_onnx_configs=models_and_onnx_configs,
-                onnx_named_outputs=onnx_outputs,
-                atol=args.atol,
-                output_dir=args.output.parent,
-                output_names=output_names,
-                device=args.device,
-                input_shapes=input_shapes,
-            )
+            models_and_onnx_configs = get_encoder_decoder_models_for_export(model, onnx_config)
+        elif task.startswith("causal-lm") and not args.monolith:
+            models_and_onnx_configs = get_decoder_models_for_export(model, onnx_config)
         else:
-            validate_model_outputs(
-                config=onnx_config,
-                reference_model=model,
-                onnx_model=args.output,
-                onnx_named_outputs=onnx_outputs,
-                atol=args.atol,
-                device=args.device,
-                input_shapes=input_shapes,
+            models_and_onnx_configs = {"model": (model, onnx_config)}
+
+    _, onnx_outputs = export_models(
+        models_and_onnx_configs=models_and_onnx_configs,
+        opset=args.opset,
+        output_dir=args.output,
+        output_names=onnx_files_subpaths,
+        input_shapes=input_shapes,
+        device=args.device,
+    )
+
+    # Optionally post process the obtained ONNX file(s), for example to merge the decoder / decoder with past if any
+    # TODO: treating stable diffusion separately is quite ugly
+    if not args.no_post_process and task != "stable-diffusion":
+        try:
+            models_and_onnx_configs, onnx_files_subpaths = onnx_config.post_process_exported_models(
+                args.output, models_and_onnx_configs, onnx_files_subpaths
+            )
+        except Exception as e:
+            raise Exception(
+                f"The post-processing of the ONNX export failed. The export can still be performed by passing the option --no-post-process. Detailed error: {e}"
             )
 
-        logger.info(f"The ONNX export succeeded and the exported model was saved at: {args.output.parent.as_posix()}")
+    try:
+        validate_models_outputs(
+            models_and_onnx_configs=models_and_onnx_configs,
+            onnx_named_outputs=onnx_outputs,
+            atol=args.atol,
+            output_dir=args.output,
+            onnx_files_subpaths=onnx_files_subpaths,
+            input_shapes=input_shapes,
+            device=args.device,
+        )
+        logger.info(f"The ONNX export succeeded and the exported model was saved at: {args.output.as_posix()}")
     except ShapeError as e:
         raise e
     except AtolError as e:
         logger.warning(
-            f"The ONNX export succeeded with the warning: {e}.\n The exported model was saved at: {args.output.parent.as_posix()}"
+            f"The ONNX export succeeded with the warning: {e}.\n The exported model was saved at: {args.output.as_posix()}"
         )
     except OutputMatchError as e:
         logger.warning(
-            f"The ONNX export succeeded with the warning: {e}.\n The exported model was saved at: {args.output.parent.as_posix()}"
+            f"The ONNX export succeeded with the warning: {e}.\n The exported model was saved at: {args.output.as_posix()}"
         )
     except Exception as e:
         logger.error(
-            f"An error occured with the error message: {e}.\n The exported model was saved at: {args.output.parent.as_posix()}"
+            f"An error occured during validation with the error message: {e}.\n The exported model was saved at: {args.output.as_posix()}"
         )
 
 

diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py
@@ -24,13 +24,19 @@
 import re
 from abc import ABC, abstractmethod
 from collections import OrderedDict
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
 
 import onnx
-from onnxruntime import InferenceSession
+from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions
 from transformers.utils import is_torch_available
 
-from ...utils import DEFAULT_DUMMY_SHAPES, DummyInputGenerator, DummyTrainingLabelsInputGenerator, logging
+from ...utils import (
+    DEFAULT_DUMMY_SHAPES,
+    DummyInputGenerator,
+    DummyTrainingLabelsInputGenerator,
+    is_diffusers_available,
+    logging,
+)
 from ...utils import TORCH_MINIMUM_VERSION as GLOBAL_MIN_TORCH_VERSION
 from ...utils.doc import add_dynamic_docstring
 from ..base import ExportConfig
@@ -41,12 +47,11 @@
 
     from transformers import PretrainedConfig, PreTrainedModel, TFPreTrainedModel
 
-
-logger = logging.get_logger(__name__)
+    if is_diffusers_available():
+        from diffusers import ModelMixin
 
 
-# 2 Gb
-EXTERNAL_DATA_FORMAT_SIZE_LIMIT = 2 * 1024 * 1024 * 1024
+logger = logging.get_logger(__name__)
 
 
 @dataclasses.dataclass
@@ -304,7 +309,9 @@ def fix_dynamic_axes(self, model_path: "Path", device: str = "cpu", input_shapes
         else:
             providers = ["CPUExecutionProvider"]
 
-        session = InferenceSession(model_path.as_posix(), providers=providers)
+        session_options = SessionOptions()
+        session_options.graph_optimization_level = GraphOptimizationLevel.ORT_DISABLE_ALL  # no need to optimize here
+        session = InferenceSession(model_path.as_posix(), providers=providers, sess_options=session_options)
 
         to_fix = []
         for output_idx, node in enumerate(session.get_outputs()):
@@ -463,14 +470,38 @@ def generate_dummy_inputs_for_validation(self, reference_model_inputs: Dict[str,
         """
         Generates inputs for ONNX Runtime using the reference model inputs. Override this to run inference with seq2seq
         models which have the encoder and decoder exported as separate ONNX files.
+
         Args:
             reference_model_inputs ([`Dict[str, Tensor]`):
                 Reference inputs for the model.
+
         Returns:
             `Dict[str, Tensor]`: The mapping holding the kwargs to provide to the model's forward function
         """
         return reference_model_inputs
 
+    def post_process_exported_models(
+        self,
+        path: "Path",
+        models_and_onnx_configs: Dict[
+            str, Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], "OnnxConfig"]
+        ],
+        onnx_files_subpaths: List[str],
+    ):
+        """
+        Performs any model-specific post-processing on the ONNX.
+
+        Args:
+            path (`Path`):
+                Path to the directory of the stored ONNX model.
+            models_and_onnx_configs (`Dict[str, Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], "OnnxConfig"]]`):
+                A dictionnary containing the models t apply post-processing on, and their corresponding ONNX configuration.
+            onnx_files_subpaths (`List[str]`):
+            The relative paths from the export directory to the ONNX files to do post-processing on. The order must be the same as*
+            the order of submodels in the ordered dict `models_and_onnx_configs`.
+        """
+        return models_and_onnx_configs, onnx_files_subpaths
+
 
 class OnnxConfigWithPast(OnnxConfig, ABC):
     """
@@ -508,6 +539,8 @@ def __init__(
                 f"use_past = {use_past} is different than use_present_in_outputs = {use_present_in_outputs}, the value "
                 "of use_present_in_outputs value will be used for the outputs."
             )
+        self.is_merged = False
+        self.use_cache_branch = None
         super().__init__(config, task=task)
 
     @classmethod
@@ -559,7 +592,11 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs):
                 if dummy_input_gen.supports_input(input_name):
                     # models from TextSeq2SeqOnnxConfig use decoder_input_ids as input name
                     # while models from TextDecoderOnnxConfig use input_ids, hence the check for both
-                    if self.use_past is True and input_name in ["decoder_input_ids", "input_ids"]:
+                    if (
+                        self.use_past is True
+                        and self.use_cache_branch is not False
+                        and input_name in ["decoder_input_ids", "input_ids"]
+                    ):
                         sequence_length = dummy_input_gen.sequence_length
                         if "sequence_length" in kwargs and kwargs["sequence_length"] != 1:
                             logger.info(
@@ -579,7 +616,12 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs):
                 )
 
         # refer to https://github.com/huggingface/optimum/pull/764
-        if self.use_past_in_inputs and "attention_mask" in dummy_inputs and self.PAD_ATTENTION_MASK_TO_PAST:
+        if (
+            self.use_past_in_inputs
+            and self.PAD_ATTENTION_MASK_TO_PAST
+            and self.use_cache_branch is not False
+            and "attention_mask" in dummy_inputs
+        ):
             past_length = dummy_inputs["past_key_values"][0][0].shape[2]
             dummy_inputs["attention_mask"] = DummyInputGenerator.pad_input_on_dim(
                 dummy_inputs["attention_mask"],
@@ -807,6 +849,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs):
         input_name, _ = next(iter(self._onnx_config.inputs.items()))
         batch_size = dummy_inputs[input_name].shape[0]
 
+        # TODO: doesn't this break attention_mask generation?
         if isinstance(self._onnx_config, OnnxSeq2SeqConfigWithPast) and self._onnx_config.use_past_in_inputs is True:
             kwargs["sequence_length"] = 1