The Tasks Manager

从一个框架导出模型到某些格式（在此也称为后端），需要指定导出函数所需的输入和输出信息。对于每个后端，optimum.exporters 的结构如下

包含每个模型执行导出所需信息的配置类。
使用适当的模型配置进行导出函数。

TasksManager 的作用是作为给定名称和任务加载模型的主要入口点，以及获取给定的（架构，后端）配对的适当配置。这样，有一个集中注册的位置来注册 task -> 模型类 和 (架构，后端) -> 配置 映射。这使得导出函数可以使用这个映射，并依赖于它提供的各种检查。

任务名称

支持的任务可能取决于后端，但以下为 PyTorch 和 TensorFlow 之间的任务名称与自动类之间的映射。

可以通过执行以下操作来了解给定后端支持哪些任务

>>> from optimum.exporters.tasks import TasksManager

>>> model_type = "distilbert"
>>> # For instance, for the ONNX export.
>>> backend = "onnx"
>>> distilbert_tasks = list(TasksManager.get_supported_tasks_for_model_type(model_type, backend).keys())

>>> print(distilbert_tasks)
['default', 'fill-mask', 'text-classification', 'multiple-choice', 'token-classification', 'question-answering']

PyTorch

任务	自动类
`text-generation`, `text-generation-with-past`	`AutoModelForCausalLM`
`feature-extraction`, `feature-extraction-with-past`	`AutoModel`
`fill-mask`	`AutoModelForMaskedLM`
`question-answering`	`AutoModelForQuestionAnswering`
`text2text-generation`, `text2text-generation-with-past`	`AutoModelForSeq2SeqLM`
`text-classification`	`AutoModelForSequenceClassification`
`token-classification`	`AutoModelForTokenClassification`
`multiple-choice`	`AutoModelForMultipleChoice`
`image-classification`	`AutoModelForImageClassification`
`object-detection`	`AutoModelForObjectDetection`
`image-segmentation`	`AutoModelForImageSegmentation`
`masked-im`	`AutoModelForMaskedImageModeling`
`semantic-segmentation`	`AutoModelForSemanticSegmentation`
`automatic-speech-recognition`	`AutoModelForSpeechSeq2Seq`

TensorFlow

任务	自动类
`text-generation`, `text-generation-with-past`	`TFAutoModelForCausalLM`
`default`, `default-with-past`	`TFAutoModel`
`fill-mask`	`TFAutoModelForMaskedLM`
`question-answering`	`TFAutoModelForQuestionAnswering`
`text2text-generation`, `text2text-generation-with-past`	`TFAutoModelForSeq2SeqLM`
`text-classification`	`TFAutoModelForSequenceClassification`
`token-classification`	`TFAutoModelForTokenClassification`
`multiple-choice`	`TFAutoModelForMultipleChoice`
`semantic-segmentation`	`TFAutoModelForSemanticSegmentation`

引用

类 optimum.exporters.TasksManager

< 源代码 >

( )

处理 任务名称 -> 模型类 和 架构 -> 配置 映射。

create_register

< source >

( backend: str overwrite_existing: bool = False ) → Callable[[str, Tuple[str, ...]], Callable[[Type], Type]]

参数

backend (str) — 注册函数将处理的后端名称。
overwrite_existing (bool, 默认 False) — 注册函数是否允许覆盖已存在的配置。

返回值

Callable[[str, Tuple[str, ...]], Callable[[Type], Type]]

接收模型类型和支持的任务的装饰器。

创建特定后端的注册函数。

示例

>>> register_for_new_backend = create_register("new-backend")

>>> @register_for_new_backend("bert", "text-classification", "token-classification")
>>> class BertNewBackendConfig(NewBackendConfig):
>>>     pass

determine_framework

< source >

( model_name_or_path: typing.Union[str, pathlib.Path] subfolder: str = '' revision: typing.Optional[str] = None cache_dir: str = '/root/.cache/huggingface/hub' token: typing.Union[bool, str, NoneType] = None ) → str

参数

model_name_or_path (Union[str, Path]) — 可以是 Hugging Face Hub 上模型仓库的模型 ID，或者包含模型的本地目录路径。
子文件夹 (str, 可选, 默认为 "") — 如果模型文件位于 Hugging Face Hub 上模型目录/库的子文件夹中，可以在此处指定子文件夹名称。
修订版 (Optional[str], 默认为 None) — 修订版是指要使用的特定模型版本。它可以是分支名称、标签名称或提交 ID。
缓存目录 (Optional[str], 可选) — 如果不使用标准缓存，请指定用于缓存的预训练模型权重的目录。
token (Optional[Union[bool,str]], 默认为 None) — 使用的token作为HTTP载体认证远程文件。如果为True，将使用在运行huggingface-cli login时生成的token（存储在 huggingface_hub.constants.HF_TOKEN_PATH）。

返回值

str

用于导出的框架。

确定用于导出的框架。

优先级按以下顺序排列

通过 framework 的用户输入。
如果提供了本地检查点，使用与检查点相同的框架。
如果模型仓库，尝试从缓存中推断框架（如果可用），否则从Hub中推断。
如果无法推断，则使用环境中可用的框架，优先使用PyTorch。

get_all_tasks

< source >

( ) → List

返回值

List

所有可能的任务。

检索所有可能的任务。

get_exporter_config_constructor

< source >

( exporter: str model: typing.Union[ForwardRef('PreTrainedModel'), ForwardRef('TFPreTrainedModel'), NoneType] = None task: str = 'feature-extraction' model_type: typing.Optional[str] = None model_name: typing.Optional[str] = None exporter_config_kwargs: typing.Union[typing.Dict[str, typing.Any], NoneType] = None library_name: typing.Optional[str] = None ) → ExportConfigConstructor

参数

exporter (str) — 要使用的导出器。
model (Optional[Union[PreTrainedModel, TFPreTrainedModel]], 默认为 None) — 模型实例。
task (str, 默认为 "feature-extraction") — 获取配置的任务。
model_type (Optional[str], 默认为 None) — 获取配置的模型类型。
model_name (Optional[str], 默认为 None) — 模型对象的name属性，仅用于异常信息。

返回值

ExportConfigConstructor

请求后端的 ExportConfig 构造函数。

获取模型（或可选的模型类型）和任务组合的 ExportConfigConstructor。

get_model_class_for_task

参数

task (str) —— 需要的任务。
framework (str, 默认为 "pt") —— 用来导出的框架。
model_type (Optional[str], 默认为 None) —— 获取模型类所需的模型类型。一些架构需要加载自定义类，而不能从自动类加载。
model_class_name (Optional[str], 默认为 None) —— 模型类名称，允许覆盖用于任务的默认类检测。此参数对于例如“自动语音识别”，它可能映射到 AutoModelForSpeechSeq2Seq 或 AutoModelForCTC 是有用的。
library (str, 默认为 transformers) — 模型的库名称。可以是“transformers”， “timm”， “diffusers”， “sentence_transformers”。

尝试从任务名称中检索 AutoModel 类。

get_model_from_task

< source >

( task: str model_name_or_path: typing.Union[str, pathlib.Path] subfolder: str = '' revision: typing.Optional[str] = None cache_dir: str = '/root/.cache/huggingface/hub' token: typing.Union[bool, str, NoneType] = None framework: typing.Optional[str] = None torch_dtype: typing.Optional[ForwardRef('torch.dtype')] = None device: typing.Union[ForwardRef('torch.device'), str, NoneType] = None library_name: typing.Optional[str] = None **model_kwargs )

参数

task (str) — 所需的任务。
model_name_or_path (Union[str, Path]) — 可以是Hugging Face Hub上模型仓库的模型ID，或者是包含模型的本地目录的路径。
subfolder (str, 默认为 ? ) — 如果模型文件位于Hugging Face Hub上的模型目录/仓库的子目录中，您可以在此指定子目录的名称。
revision (Optional[str], 可选) — 修订版是使用的特定模型版本。它可以是一个分支名称、一个标签名称或一个提交ID。
cache_dir (Optional[str], 可选) — 如果不使用标准缓存，已缓存下载的预训练模型权重的目录路径。
token (Optional[Union[bool,str]], 默认为 None) — 用于远程文件的HTTP bearer授权的令牌。如果为 True，将使用运行 huggingface-cli login 时生成的令牌（存储在 huggingface_hub.constants.HF_TOKEN_PATH 中）。
framework (Optional[str], 可选) — 导出时使用的框架。如果没有提供，请参考 TasksManager.determine_framework 获取优先级。
torch_dtype (Optional[torch.dtype], defaults to None) — 加载模型的数值类型。仅支持PyTorch。
device (Optional[torch.device], defaults to None) — 在该设备上初始化模型。仅支持PyTorch。对于PyTorch，默认为“cpu”。
library_name (Optional[str], defaults to None) — 模型的库名。可以是“transformers”、“timm”、“diffusers”、“sentence_transformers”之一。若未提供，请查阅TasksManager.infer_library_from_model以获取优先级。
model_kwargs (Dict[str, Any], 可选) — 向模型的 .from_pretrained() 方法传递的关键字参数。

从名称和要启用的任务中检索模型。

get_supported_model_type_for_task

< 源码 >

( task: str exporter: str )

返回对于给定任务的导出器支持的反转架构列表。Transformers特定。

get_supported_tasks_for_model_type

< 源码 >

( model_type: str exporter: str model_name: typing(Optional[str]) = None library_name: typing(Optional[str]) = None ) → TaskNameToExportConfigDict

参数

model_type (str) — 获取支持任务的模型类型。
exporter (str) — 导出器的名称。
model_name (Optional[str], 默认 None) — 模型对象的名称属性，仅用于异常信息。
library_name (Optional[str], 默认为 None) — 模型库的名称。可以是以下之一：“transformers”, “timm”, “diffusers”, “sentence_transformers”。

返回值

TaskNameToExportConfigDict

将每个任务映射到相应的 ExportConfig 构造函数的字典。

从模型类型中检索 task -> exporter backend config constructors 映射。

infer_library_from_model

< 源代码 >

( model: typing.Union[str, ForwardRef('PreTrainedModel'), ForwardRef('TFPreTrainedModel'), ForwardRef('DiffusionPipeline'), typing.Type] subfolder: str = '' revision: typing.Optional[str] = None cache_dir: str = '/root/.cache/huggingface/hub' token: typing.Union[bool, str, NoneType] = None ) → str

参数

model (str, 预训练模型实例, TF预训练模型, 扩散流水线, 类型) — 从中推断任务的模型。这可以是HuggingFace Hub上的存储库名称，一个模型实例，或一个模型类。
subfolder (str，默认为""） — 如果模型文件位于HuggingFace Hub上的模型目录/存储库的子目录中，您可以在此处指定子目录名称。
revision (Optional[str]，可选，默认为None） — 修订版本是使用的特定模型版本。它可以是一个分支名称、一个标记名称或一个提交ID。
cache_dir (Optional[str], 可选) — 如果不使用标准缓存，则缓存预训练模型权重的目录路径。
token (Optional[Union[bool,str]], 默认为 None) — 用作远程文件HTTP携带授权的token。如果设置为 True，将使用在运行 huggingface-cli login 时生成的token（存储在 huggingface_hub.constants.HF_TOKEN_PATH）。

返回值

str

自动从模型存储库、模型实例或模型类中检测到的库名称。

从模型存储库、模型实例或模型类中推断库。

infer_task_from_model

< 源代码 >

参数

model (Union[str, PreTrainedModel, TFPreTrainedModel, DiffusionPipeline, Type]) — 从中推断任务的模型。这可以是HuggingFace Hub上存储库的名称、模型的一个实例或模型类。
subfolder (str, 可选的，默认为"") — 如果模型文件位于HuggingFace Hub上模型目录/存储库的子目录中，您可以在此指定子目录名称。
revision (Optional[str], 默认为None) — 修订版本是特定模型版本。它可以是一个分支名称、标签名称或提交ID。
cache_dir (Optional[str], optional) — 如果不使用标准缓存，则缓存已下载预训练模型权重的目录路径。
token (Optional[Union[bool,str]], 默认值为None) — 作为远程文件HTTP承载授权使用的令牌。如果为True，将使用当运行huggingface-cli login时生成的令牌（存储在huggingface_hub.constants.HF_TOKEN_PATH）。

返回值

str

从HF hub仓库、模型实例或模型类自动检测到的任务名称。

从模型仓库、模型实例或模型类推断任务。

standardize_model_attributes

< source >

( model: typing.Union[ForwardRef('PreTrainedModel'), ForwardRef('TFPreTrainedModel'), ForwardRef('DiffusionPipeline')] )

参数

模型 (Union[PreTrainedModel, TFPreTrainedModel, DiffusionPipeline]) — 模型的实例。

更新模型以供导出。此函数适用于对来自不同库的模型进行必须更改，以遵循transformer风格。

< > 在GitHub上更新

最佳性能优化工具

The Tasks Manager

任务名称

PyTorch

TensorFlow

引用

类 optimum.exporters.TasksManager

create_register

determine_framework

get_all_tasks

get_exporter_config_constructor

get_model_class_for_task

get_model_from_task

get_supported_model_type_for_task

get_supported_tasks_for_model_type

infer_library_from_model

infer_task_from_model

standardize_model_attributes