refactoring

optimum/commands/__init__.py

@@ -12,20 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from transformers.utils import is_tf_available
-
-from ..utils import is_onnx_available, is_onnxruntime_available
 from .base import BaseOptimumCLICommand, CommandInfo, RootOptimumCLICommand
 from .env import EnvironmentCommand
-from .export import ExportCommand
+from .export import ExportCommand, ONNXExportCommand, TFLiteExportCommand
+from .onnxruntime import ONNXRuntimeCommand, ONNXRuntimmeOptimizeCommand, ONNXRuntimmeQuantizeCommand
 from .optimum_cli import register_optimum_cli_subcommand
-
-
-if is_onnx_available():
-    from .export import ONNXExportCommand
-
-if is_tf_available():
-    from .export import TFLiteExportCommand
-
-if is_onnxruntime_available():
-    from .onnxruntime import ONNXRuntimeCommand, ONNXRuntimmeOptimizeCommand, ONNXRuntimmeQuantizeCommand

optimum/commands/base.py

@@ -118,6 +118,9 @@ def registered_subcommands(self):
             self._registered_subcommands = []
         return self._registered_subcommands
 
+    def check_requirements(self):
+        pass
+
     @staticmethod
     def parse_args(parser: "ArgumentParser"):
         pass

optimum/commands/export/__init__.py

@@ -12,13 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from transformers.utils import is_tf_available
-
-from ...utils import is_onnx_available
 from .base import ExportCommand
-
-
-if is_onnx_available():
-    from .onnx import ONNXExportCommand
-if is_tf_available():
-    from .tflite import TFLiteExportCommand
+from .onnx import ONNXExportCommand
+from .tflite import TFLiteExportCommand

optimum/commands/export/base.py

@@ -14,36 +14,25 @@
 # limitations under the License.
 """optimum.exporters command-line interface base classes."""
 
-from transformers.utils import is_tf_available
-
-from ...utils import is_onnx_available
 from .. import BaseOptimumCLICommand, CommandInfo
+from .onnx import ONNXExportCommand
+from .tflite import TFLiteExportCommand
 
 
 class ExportCommand(BaseOptimumCLICommand):
     COMMAND = CommandInfo(
         name="export",
         help="Export PyTorch and TensorFlow models to several format.",
     )
-    SUBCOMMANDS = ()
-    if is_onnx_available():
-        from .onnx import ONNXExportCommand
-
-        SUBCOMMANDS += (
-            CommandInfo(
-                name="onnx",
-                help="Export PyTorch and TensorFlow to ONNX.",
-                subcommand_class=ONNXExportCommand,
-            ),
-        )
-
-    if is_tf_available():
-        from .tflite import TFLiteExportCommand
-
-        SUBCOMMANDS += (
-            CommandInfo(
-                name="tflite",
-                help="Export TensorFlow to TensorFlow Lite.",
-                subcommand_class=TFLiteExportCommand,
-            ),
-        )
+    SUBCOMMANDS = (
+        CommandInfo(
+            name="onnx",
+            help="Export PyTorch and TensorFlow to ONNX.",
+            subcommand_class=ONNXExportCommand,
+        ),
+        CommandInfo(
+            name="tflite",
+            help="Export TensorFlow to TensorFlow Lite.",
+            subcommand_class=TFLiteExportCommand,
+        ),
+    )

optimum/commands/export/onnx.py

@@ -13,29 +13,200 @@
 # limitations under the License.
 """Defines the command line for the export with ONNX."""
 
+import argparse
+from pathlib import Path
 from typing import TYPE_CHECKING
 
-from ...exporters.onnx.__main__ import main_export, parse_args_onnx
-from ...utils import DEFAULT_DUMMY_SHAPES
+from ...exporters import TasksManager
+from ...utils import DEFAULT_DUMMY_SHAPES, is_onnx_available
 from ..base import BaseOptimumCLICommand
 
 
 if TYPE_CHECKING:
     from argparse import ArgumentParser
 
 
+def parse_args_onnx(parser):
+    required_group = parser.add_argument_group("Required arguments")
+    required_group.add_argument(
+        "-m", "--model", type=str, required=True, help="Model ID on huggingface.co or path on disk to load model from."
+    )
+    required_group.add_argument(
+        "output", type=Path, help="Path indicating the directory where to store the generated ONNX model."
+    )
+
+    optional_group = parser.add_argument_group("Optional arguments")
+    optional_group.add_argument(
+        "--task",
+        default="auto",
+        help=(
+            "The task to export the model for. If not specified, the task will be auto-inferred based on the model. Available tasks depend on the model, but are among:"
+            f" {str(list(TasksManager._TASKS_TO_AUTOMODELS.keys()))}. For decoder models, use `xxx-with-past` to export the model using past key values in the decoder."
+        ),
+    )
+    optional_group.add_argument(
+        "--opset",
+        type=int,
+        default=None,
+        help="If specified, ONNX opset version to export the model with. Otherwise, the default opset for the given model architecture will be used.",
+    )
+    optional_group.add_argument(
+        "--device",
+        type=str,
+        default="cpu",
+        help='The device to use to do the export. Defaults to "cpu".',
+    )
+    optional_group.add_argument(
+        "--fp16",
+        action="store_true",
+        help="Use half precision during the export. PyTorch-only, requires `--device cuda`.",
+    )
+    optional_group.add_argument(
+        "--optimize",
+        type=str,
+        default=None,
+        choices=["O1", "O2", "O3", "O4"],
+        help=(
+            "Allows to run ONNX Runtime optimizations directly during the export. Some of these optimizations are specific to ONNX Runtime, and the resulting ONNX will not be usable with other runtime as OpenVINO or TensorRT. Possible options:\n"
+            "    - O1: Basic general optimizations\n"
+            "    - O2: Basic and extended general optimizations, transformers-specific fusions\n"
+            "    - O3: Same as O2 with GELU approximation\n"
+            "    - O4: Same as O3 with mixed precision (fp16, GPU-only, requires `--device cuda`)"
+        ),
+    )
+    optional_group.add_argument(
+        "--monolith",
+        action="store_true",
+        help=(
+            "Forces to export the model as a single ONNX file. By default, the ONNX exporter may break the model in several"
+            " ONNX files, for example for encoder-decoder models where the encoder should be run only once while the"
+            " decoder is looped over."
+        ),
+    )
+    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."
+        ),
+    )
+    optional_group.add_argument(
+        "--framework",
+        type=str,
+        choices=["pt", "tf"],
+        default=None,
+        help=(
+            "The framework to use for the ONNX export."
+            " If not provided, will attempt to use the local checkpoint's original framework"
+            " or what is available in the environment."
+        ),
+    )
+    optional_group.add_argument(
+        "--atol",
+        type=float,
+        default=None,
+        help="If specified, the absolute difference tolerance when validating the model. Otherwise, the default atol for the model will be used.",
+    )
+    optional_group.add_argument("--cache_dir", type=str, default=None, help="Path indicating where to store cache.")
+    optional_group.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        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(
+        "--pad_token_id",
+        type=int,
+        default=None,
+        help=(
+            "This is needed by some models, for some tasks. If not provided, will attempt to use the tokenizer to guess"
+            " it."
+        ),
+    )
+
+    input_group = parser.add_argument_group(
+        "Input shapes (if necessary, this allows to override the shapes of the input given to the ONNX exporter, that requires an example input)."
+    )
+    doc_input = "to use in the example input given to the ONNX export."
+    input_group.add_argument(
+        "--batch_size",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["batch_size"],
+        help=f"Text tasks only. Batch size {doc_input}",
+    )
+    input_group.add_argument(
+        "--sequence_length",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["sequence_length"],
+        help=f"Text tasks only. Sequence length {doc_input}",
+    )
+    input_group.add_argument(
+        "--num_choices",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["num_choices"],
+        help=f"Text tasks only. Num choices {doc_input}",
+    )
+    input_group.add_argument(
+        "--width",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["width"],
+        help=f"Image tasks only. Width {doc_input}",
+    )
+    input_group.add_argument(
+        "--height",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["height"],
+        help=f"Image tasks only. Height {doc_input}",
+    )
+    input_group.add_argument(
+        "--num_channels",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["num_channels"],
+        help=f"Image tasks only. Number of channels {doc_input}",
+    )
+    input_group.add_argument(
+        "--feature_size",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["feature_size"],
+        help=f"Audio tasks only. Feature size {doc_input}",
+    )
+    input_group.add_argument(
+        "--nb_max_frames",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["nb_max_frames"],
+        help=f"Audio tasks only. Maximum number of frames {doc_input}",
+    )
+    input_group.add_argument(
+        "--audio_sequence_length",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["audio_sequence_length"],
+        help=f"Audio tasks only. Audio sequence length {doc_input}",
+    )
+
+    # deprecated argument
+    parser.add_argument("--for-ort", action="store_true", help=argparse.SUPPRESS)
+
+
 class ONNXExportCommand(BaseOptimumCLICommand):
     @staticmethod
     def parse_args(parser: "ArgumentParser"):
         return parse_args_onnx(parser)
 
+    def check_requirements(self):
+        if is_onnx_available():
+            from ...exporters.onnx.__main__ import main_export  # noqa
+        else:
+            raise ImportError("Onnx is not installed. Please install Onnx first.")
+
     def run(self):
+        self.check_requirements()
+
         # Get the shapes to be used to generate dummy inputs
         input_shapes = {}
         for input_name in DEFAULT_DUMMY_SHAPES.keys():
             input_shapes[input_name] = getattr(self.args, input_name)
 
-        main_export(
+        main_export(  # noqa: F821
             model_name_or_path=self.args.model,
             output=self.args.output,
             task=self.args.task,

optimum/commands/export/tflite.py

@@ -18,6 +18,8 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional
 
+from transformers.utils import is_tf_available
+
 from ...exporters import TasksManager
 from ...exporters.tflite import QuantizationApproach
 from ..base import BaseOptimumCLICommand
@@ -237,6 +239,11 @@ def __init__(
     def parse_args(parser: "ArgumentParser"):
         return parse_args_tflite(parser)
 
+    def check_requirements(self):
+        if not is_tf_available():
+            raise ImportError("TensorFlow is not installed. Please install TensorFlow first.")
+
     def run(self):
+        self.check_requirements()
         full_command = f"python3 -m optimum.exporters.tflite {self.args_string}"
         subprocess.run(full_command, shell=True, check=True)

optimum/commands/onnxruntime/optimize.py

@@ -19,8 +19,7 @@
 
 from optimum.commands.base import BaseOptimumCLICommand
 
-from ...onnxruntime.configuration import AutoOptimizationConfig, ORTConfig
-from ...onnxruntime.optimization import ORTOptimizer
+from ...utils import is_onnxruntime_available
 
 
 if TYPE_CHECKING:
@@ -77,26 +76,35 @@ class ONNXRuntimmeOptimizeCommand(BaseOptimumCLICommand):
     def parse_args(parser: "ArgumentParser"):
         return parse_args_onnxruntime_optimize(parser)
 
+    def check_requirements(self):
+        if is_onnxruntime_available():
+            from ...onnxruntime.configuration import AutoOptimizationConfig, ORTConfig  # noqa
+            from ...onnxruntime.optimization import ORTOptimizer  # noqa
+        else:
+            raise ImportError("Onnxruntime is not installed. Please install Onnxruntime first.")
+
     def run(self):
+        self.check_requirements()
+
         if self.args.output == self.args.onnx_model:
             raise ValueError("The output directory must be different than the directory hosting the ONNX model.")
 
         save_dir = self.args.output
 
         file_names = [model.name for model in self.args.onnx_model.glob("*.onnx")]
-        optimizer = ORTOptimizer.from_pretrained(self.args.onnx_model, file_names)
+        optimizer = ORTOptimizer.from_pretrained(self.args.onnx_model, file_names)  # noqa: F821
 
         if self.args.config:
-            optimization_config = ORTConfig
+            optimization_config = ORTConfig  # noqa: F821
         elif self.args.O1:
-            optimization_config = AutoOptimizationConfig.O1()
+            optimization_config = AutoOptimizationConfig.O1()  # noqa: F821
         elif self.args.O2:
-            optimization_config = AutoOptimizationConfig.O2()
+            optimization_config = AutoOptimizationConfig.O2()  # noqa: F821
         elif self.args.O3:
-            optimization_config = AutoOptimizationConfig.O3()
+            optimization_config = AutoOptimizationConfig.O3()  # noqa: F821
         elif self.args.O4:
-            optimization_config = AutoOptimizationConfig.O4()
+            optimization_config = AutoOptimizationConfig.O4()  # noqa: F821
         else:
-            optimization_config = ORTConfig.from_pretained(self.args.config).optimization
+            optimization_config = ORTConfig.from_pretained(self.args.config).optimization  # noqa: F821
 
         optimizer.optimize(save_dir=save_dir, optimization_config=optimization_config)