Optimum 文档

导出函数

您正在查看的是需要从源码安装。如果您希望进行常规的 pip 安装,请查看最新的稳定版本 (v1.27.0)。
Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

导出函数

您可以从 🤗 Optimum 中的两个框架将模型导出到 ONNX:PyTorch 和 TensorFlow。每个框架都有一个导出函数,export_pytorch()export_tensorflow(),但推荐的使用方法是通过主导出函数 ~optimum.exporters.main_export,它将根据可用的框架负责使用适当的导出函数,检查导出的模型是否有效,并提供扩展选项以对导出的模型运行优化。

主要功能

optimum.exporters.onnx.main_export

< >

( model_name_or_path: str output: typing.Union[str, pathlib.Path] task: str = 'auto' opset: typing.Optional[int] = None device: str = 'cpu' dtype: typing.Optional[str] = None fp16: typing.Optional[bool] = False optimize: typing.Optional[str] = None monolith: bool = False no_post_process: bool = False framework: typing.Optional[str] = None atol: typing.Optional[float] = None pad_token_id: typing.Optional[int] = None subfolder: str = '' revision: str = 'main' force_download: bool = False local_files_only: bool = False trust_remote_code: bool = False cache_dir: str = '/root/.cache/huggingface/hub' token: typing.Union[str, bool, NoneType] = None for_ort: bool = False do_validation: bool = True model_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None custom_onnx_configs: typing.Optional[typing.Dict[str, ForwardRef('OnnxConfig')]] = None fn_get_submodels: typing.Optional[typing.Callable] = None use_subprocess: bool = False _variant: str = 'default' library_name: typing.Optional[str] = None legacy: bool = False no_dynamic_axes: bool = False do_constant_folding: bool = True slim: bool = False **kwargs_shapes )

必填参数

  • model_name_or_path (str) — Huggingface.co 上的模型 ID 或磁盘上要导出的模型存储库路径。示例:model_name_or_path="BAAI/bge-m3"mode_name_or_path="/path/to/model_folder
  • output (Union[str, Path]) — 指示存储生成的 ONNX 模型的目录的路径。

可选参数

  • task (Optional[str], 默认为 None) — 导出模型所针对的任务。如果未指定,将根据模型自动推断任务。对于解码器模型,使用 xxx-with-past 导出模型时,解码器中使用过去的键值。
  • opset (Optional[int], 默认为 None) — 如果指定,则为导出模型使用的 ONNX opset 版本。否则,将使用给定模型架构的默认 opset。
  • device (str, 默认为 "cpu") — 用于执行导出的设备。默认为“cpu”。
  • fp16 (Optional[bool], 默认为 "False") — 导出时使用半精度。仅限 PyTorch,需要 device="cuda"
  • dtype (Optional[str], 默认为 None) — 导出时使用的浮点精度。支持的选项:"fp32" (float32)、"fp16" (float16)、"bf16" (bfloat16)。默认为 "fp32"
  • optimize (Optional[str], 默认为 None) — 允许在导出期间直接运行 ONNX Runtime 优化。其中一些优化是 ONNX Runtime 特有的,因此生成的 ONNX 模型无法与其他运行时(如 OpenVINO 或 TensorRT)一起使用。可用选项:"O1", "O2", "O3", "O4"。参考:AutoOptimizationConfig
  • monolith (bool, 默认为 False) — 强制将模型导出为单个 ONNX 文件。
  • no_post_process (bool, 默认为 False) — 允许禁用对导出的 ONNX 模型执行的任何默认后处理。
  • framework (Optional[str], 默认为 None) — 用于 ONNX 导出的框架("pt""tf")。如果未提供,将尝试自动检测检查点的框架。
  • atol (Optional[float], 默认为 None) — 如果指定,则为验证模型时的绝对差容差。否则,将使用模型的默认 atol。
  • cache_dir (Optional[str], 默认为 None) — 指示存储缓存的路径。默认情况下将使用 Hugging Face 默认缓存路径。
  • trust_remote_code (bool, 默认为 False) — 允许使用模型存储库中托管的自定义代码进行建模。此选项仅应在您信任且已阅读代码的存储库中设置,因为它将在您的本地机器上执行模型存储库中存在的任意代码。
  • pad_token_id (Optional[int], 默认为 None) — 某些模型在某些任务中需要此项。如果未提供,将尝试使用分词器猜测它。
  • subfolder (str, 默认为 "") — 如果相关文件位于模型仓库的子文件夹中(本地或 huggingface.co 上),您可以在此处指定文件夹名称。
  • revision (str, 默认为 "main") — 修订版本是要使用的特定模型版本。它可以是分支名称、标签名称或提交 ID。
  • force_download (bool, 默认为 False) — 是否强制(重新)下载模型权重和配置文件,如果已存在则覆盖缓存版本。
  • local_files_only (Optional[bool], 默认为 False) — 是否仅查看本地文件(即不尝试下载模型)。
  • use_auth_token (Optional[Union[bool,str]], 默认为 None) — 已弃用。请改用 token 参数。
  • token (Optional[Union[bool,str]], 默认为 None) — 用作远程文件 HTTP 持票人授权的令牌。如果为 True,将使用运行 huggingface-cli login 时生成的令牌(存储在 huggingface_hub.constants.HF_TOKEN_PATH 中)。
  • model_kwargs (Optional[Dict[str, Any]], 默认为 None) — 实验性用法:在导出期间传递给模型的关键字参数。此参数应与 custom_onnx_configs 参数一起使用,例如,当模型输入/输出发生更改时(例如,如果传递了 model_kwargs={"output_attentions": True})。
  • custom_onnx_configs (Optional[Dict[str, OnnxConfig]], 默认为 None) — 实验性用法:覆盖用于给定模型的默认 ONNX 配置。此参数对于希望对导出进行更精细控制的高级用户可能很有用。示例请参见此处
  • fn_get_submodels (Optional[Callable], 默认为 None) — 实验性用法:覆盖导出时使用的默认子模型。当导出需要拆分 ONNX 的自定义架构时(例如编码器-解码器),这尤其有用。如果未指定自定义模型,optimum 将尝试使用给定任务的默认子模型,但不保证成功。
  • use_subprocess (bool, 默认为 False) — 在子进程中进行 ONNX 导出模型验证。当在 CUDA 设备上导出时,这尤其有用,因为 ORT 在推理会话销毁时不会释放内存。当设置为 True 时,main_export 调用应受 if __name__ == "__main__": 块的保护。
  • _variant (str, 默认为 default) — 指定要使用的 ONNX 导出变体。
  • library_name (Optional[str], 默认为 None) — 模型的库("transformers""diffusers""timm""sentence_transformers")。如果未提供,将尝试自动检测检查点的库名称。
  • legacy (bool, 默认为 False) — 禁用文本生成模型使用 position_ids 进行批处理生成。此外,允许将仅解码器模型导出为三个文件(不带 past、带 past 和合并模型)。此参数是为向后兼容性而引入的,将在 Optimum 的未来版本中移除。
  • no_dynamic_axes (bool, 默认为 False) — 如果为 True,则在 ONNX 导出期间禁用动态轴的使用。
  • do_constant_folding (bool, 默认为 True) — PyTorch 特有参数。如果为 True,PyTorch ONNX 导出将尽可能将常量折叠到相邻节点中。
  • slim (bool, 默认为 False) — PyTorch 特有参数。如果为 True,则使用 onnxslim 优化 ONNX 模型。
  • **kwargs_shapes (Dict) — 推理期间使用的形状。此参数允许覆盖 ONNX 导出期间使用的默认形状。

全套 ONNX 导出函数,从 Hugging Face Hub 上的模型 ID 或本地模型存储库导出。

用法示例

>>> from optimum.exporters.onnx import main_export

>>> main_export("gpt2", output="gpt2_onnx/")

optimum.exporters.onnx.onnx_export_from_model

< >

( model: typing.Union[ForwardRef('PreTrainedModel'), ForwardRef('TFPreTrainedModel'), ForwardRef('DiffusionPipeline')] output: typing.Union[str, pathlib.Path] opset: typing.Optional[int] = None optimize: typing.Optional[str] = None monolith: bool = False no_post_process: bool = False atol: typing.Optional[float] = None do_validation: bool = True model_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None custom_onnx_configs: typing.Optional[typing.Dict[str, ForwardRef('OnnxConfig')]] = None fn_get_submodels: typing.Optional[typing.Callable] = None _variant: str = 'default' legacy: bool = False preprocessors: typing.List = None device: str = 'cpu' no_dynamic_axes: bool = False task: typing.Optional[str] = None use_subprocess: bool = False do_constant_folding: bool = True slim: bool = False **kwargs_shapes )

必填参数

  • model (Union["PreTrainedModel", "TFPreTrainedModel"]) — 要导出为 ONNX 的 PyTorch 或 TensorFlow 模型。
  • output (Union[str, Path]) — 指示存储生成的 ONNX 模型的目录的路径。

可选参数

  • task (Optional[str], 默认为 None) — 导出模型所针对的任务。如果未指定,将根据模型自动推断任务。
  • opset (Optional[int], 默认为 None) — 如果指定,则为导出模型使用的 ONNX opset 版本。否则,将使用给定模型架构的默认 opset。
  • device (str, 默认为 "cpu") — 用于执行导出的设备。默认为“cpu”。
  • optimize (Optional[str], 默认为 None) — 允许在导出期间直接运行 ONNX Runtime 优化。其中一些优化是 ONNX Runtime 特有的,因此生成的 ONNX 模型无法与其他运行时(如 OpenVINO 或 TensorRT)一起使用。可用选项:"O1", "O2", "O3", "O4"。参考:AutoOptimizationConfig
  • monolith (bool, 默认为 False) — 强制将模型导出为单个 ONNX 文件。
  • no_post_process (bool, 默认为 False) — 允许禁用对导出的 ONNX 模型执行的任何默认后处理。
  • atol (Optional[float], 默认为 None) — 如果指定,则为验证模型时的绝对差容差。否则,将使用模型的默认 atol。
  • model_kwargs (Optional[Dict[str, Any]], 默认为 None) — 实验性用法:在导出期间传递给模型的关键字参数。此参数应与 custom_onnx_configs 参数一起使用,例如,当模型输入/输出发生更改时(例如,如果传递了 model_kwargs={"output_attentions": True})。
  • custom_onnx_configs (Optional[Dict[str, OnnxConfig]], 默认为 None) — 实验性用法:覆盖用于给定模型的默认 ONNX 配置。此参数对于希望对导出进行更精细控制的高级用户可能很有用。示例请参见此处
  • fn_get_submodels (Optional[Callable], 默认为 None) — 实验性用法:覆盖导出时使用的默认子模型。当导出需要拆分 ONNX 的自定义架构时(例如编码器-解码器),这尤其有用。如果未指定自定义模型,optimum 将尝试使用给定任务的默认子模型,但不保证成功。
  • use_subprocess (bool, 默认为 False) — 在子进程中执行 ONNX 导出模型的验证。当在 CUDA 设备上导出时,这尤其有用,因为 ORT 在推理会话销毁时不会释放内存。当设置为 True 时,main_export 调用应在 if __name__ == "__main__": 块中进行保护。
  • _variant (str, 默认为 default) — 指定要使用的 ONNX 导出变体。
  • legacy (bool, 默认为 False) — 禁用文本生成模型使用 position_ids 进行批处理生成。此外,允许将仅解码器模型导出为三个文件(不带 past、带 past 和合并模型)。此参数是为向后兼容性而引入的,将在 Optimum 的未来版本中移除。
  • no_dynamic_axes (bool, defaults to False) — 如果为 True,在 ONNX 导出期间禁用动态轴的使用。
  • do_constant_folding (bool, defaults to True) — PyTorch 特有参数。如果为 True,PyTorch ONNX 导出将尽可能地将常量折叠到相邻节点中。
  • slim (bool, defaults to False) — 使用 onnxslim 优化 ONNX 模型。
  • **kwargs_shapes (Dict) — 用于推理的形状。此参数允许覆盖 ONNX 导出期间使用的默认形状。

全套 ONNX 导出函数,**从预加载的 PyTorch 或 Tensorflow 模型**导出。此函数在需要对模型进行修改(例如,在导出到 ONNX 之前覆盖前向调用)的情况下特别有用。

用法示例

>>> from transformers import AutoModelForCausalLM

>>> model = AutoModelForCausalLM.from_pretrained("gpt2")
>>> # At this point, we could override some submodules, forward methods, weights, etc. from the model.

>>> onnx_export_from_model(model, output="gpt2_onnx/")

optimum.exporters.onnx.export

< >

( model: typing.Union[ForwardRef('PreTrainedModel'), ForwardRef('TFPreTrainedModel'), ForwardRef('ModelMixin')] config: OnnxConfig output: Path opset: typing.Optional[int] = None device: str = 'cpu' input_shapes: typing.Optional[typing.Dict] = None disable_dynamic_axes_fix: typing.Optional[bool] = False dtype: typing.Optional[str] = None no_dynamic_axes: bool = False do_constant_folding: bool = True model_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None ) Tuple[List[str], List[str]]

参数

  • model (PreTrainedModelTFPreTrainedModel) — 要导出的模型。
  • config (OnnxConfig) — 与导出模型关联的 ONNX 配置。
  • output (Path) — 存储导出 ONNX 模型的目录。
  • opset (Optional[int], defaults to None) — 要使用的 ONNX 运算符集的版本。
  • device (Optional[str], defaults to "cpu") — ONNX 模型将导出的设备。可以是 cpucuda。CUDA 设备上的导出仅支持 PyTorch。
  • input_shapes (Optional[Dict], defaults to None) — 如果指定,允许为提供给 ONNX 导出器的示例输入使用特定形状。
  • disable_dynamic_axes_fix (Optional[bool], defaults to False) — 是否禁用默认的动态轴修复。
  • dtype (Optional[str], defaults to None) — 将模型输入重新映射到的数据类型。仅限 PyTorch。仅支持 fp16
  • no_dynamic_axes (bool, defaults to False) — 如果为 True,在 ONNX 导出期间禁用动态轴的使用。
  • do_constant_folding (bool, defaults to True) — PyTorch 特有参数。如果为 True,PyTorch ONNX 导出将尽可能地将常量折叠到相邻节点中。
  • model_kwargs (Optional[Dict[str, Any]], defaults to None) — 实验性用法:在导出期间传递给模型的关键字参数。此参数应与 custom_onnx_config 参数一起使用,例如,当模型输入/输出发生更改时(例如,如果传递了 model_kwargs={"output_attentions": True})。

返回

Tuple[List[str], List[str]]

一个元组,包含模型输入的有序列表和 ONNX 配置中的命名输出。

将 PyTorch 或 TensorFlow 模型导出到 ONNX 中间表示。

optimum.exporters.onnx.convert.export_pytorch

< >

( model: typing.Union[ForwardRef('PreTrainedModel'), ForwardRef('ModelMixin')] config: OnnxConfig opset: int output: Path device: str = 'cpu' input_shapes: typing.Optional[typing.Dict] = None no_dynamic_axes: bool = False do_constant_folding: bool = True model_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None ) Tuple[List[str], List[str]]

参数

  • model (PreTrainedModel) — 要导出的模型。
  • config (OnnxConfig) — 与导出模型关联的 ONNX 配置。
  • opset (int) — 要使用的 ONNX 运算符集的版本。
  • output (Path) — 导出 ONNX 文件的保存路径。
  • device (str, defaults to "cpu") — ONNX 模型将导出的设备。可以是 cpucuda。CUDA 设备上的导出仅支持 PyTorch。
  • input_shapes (Optional[Dict], defaults to None) — 如果指定,允许为提供给 ONNX 导出器的示例输入使用特定形状。
  • no_dynamic_axes (bool, defaults to False) — 如果为 True,在 ONNX 导出期间禁用动态轴的使用。
  • do_constant_folding (bool, defaults to True) — PyTorch 特有参数。如果为 True,PyTorch ONNX 导出将尽可能地将常量折叠到相邻节点中。
  • model_kwargs (Optional[Dict[str, Any]], defaults to None) — 实验性用法:在导出期间传递给模型的关键字参数。此参数应与 custom_onnx_config 参数一起使用,例如,当模型输入/输出发生更改时(例如,如果传递了 model_kwargs={"output_attentions": True})。

返回

Tuple[List[str], List[str]]

一个元组,包含模型输入的有序列表和 ONNX 配置中的命名输出。

将 PyTorch 模型导出到 ONNX 中间表示。

optimum.exporters.onnx.convert.export_tensorflow

< >

( model: TFPreTrainedModel config: OnnxConfig opset: int output: Path ) Tuple[List[str], List[str]]

参数

  • model (TFPreTrainedModel) — 要导出的模型。
  • config (OnnxConfig) — 与导出模型关联的 ONNX 配置。
  • opset (int) — 要使用的 ONNX 运算符集的版本。
  • output (Path) — 存储导出 ONNX 模型的目录。
  • device (Optional[str], defaults to "cpu") — ONNX 模型将导出的设备。可以是 cpucuda。CUDA 设备上的导出仅支持 PyTorch。

返回

Tuple[List[str], List[str]]

一个元组,包含模型输入的有序列表和 ONNX 配置中的命名输出。

将 TensorFlow 模型导出到 ONNX 中间表示。

实用函数

optimum.exporters.utils.check_dummy_inputs_are_allowed

< >

( model: typing.Union[ForwardRef('PreTrainedModel'), ForwardRef('TFPreTrainedModel'), ForwardRef('ModelMixin')] dummy_input_names: typing.Iterable[str] )

参数

  • model (Union[transformers.PreTrainedModel, transformers.TFPreTrainedModel]) — 模型实例。
  • model_inputs (Iterable[str]) — 模型输入名称。

检查 ONNX 配置中的虚拟输入是否是 model 允许输入的子集。

optimum.exporters.onnx.validate_model_outputs

< >

( config: OnnxConfig reference_model: typing.Union[ForwardRef('PreTrainedModel'), ForwardRef('TFPreTrainedModel'), ForwardRef('ModelMixin')] onnx_model: Path onnx_named_outputs: typing.List[str] atol: typing.Optional[float] = None input_shapes: typing.Optional[typing.Dict] = None device: str = 'cpu' use_subprocess: typing.Optional[bool] = True model_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None )

参数

  • config (~OnnxConfig — 用于导出模型的配置。
  • reference_model (~PreTrainedModel~TFPreTrainedModel) — 用于导出的模型。
  • onnx_model (Path) — 导出模型的路径。
  • onnx_named_outputs (List[str]) — 要检查的输出名称。
  • atol (Optional[float], defaults to None) — 参考模型和导出模型之间输出差异的绝对容差。
  • input_shapes (Optional[Dict], defaults to None) — 如果指定,允许使用特定形状来验证 ONNX 模型。
  • device (str, defaults to "cpu") — ONNX 模型将验证的设备。可以是 cpucuda。CUDA 设备上的验证仅支持 PyTorch。
  • use_subprocess (Optional[bool], defaults to True) — 在子进程中启动每个导出模型的验证。
  • model_kwargs (Optional[Dict[str, Any]], defaults to None) — 实验性用法:在导出和验证期间传递给模型的关键字参数。

引发

ValueError

  • ValueError — 如果参考模型和导出模型之间的输出形状或值不匹配。

通过检查参考模型和导出模型之间的输出是否匹配来验证导出。

< > 在 GitHub 上更新