任务管理器

将模型从一个框架导出到某种格式（此处也称为后端）需要指定导出功能所需的输入和输出信息。`optimum.exporters` 为每个后端构建的方式如下：

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

TasksManager 的作用是作为主要入口点，根据名称和任务加载模型，并获取给定（架构、后端）组合的适当配置。这样，就有一个集中的位置来注册 `任务 -> 模型类` 和 `（架构、后端）-> 配置` 映射。这允许导出功能利用此功能，并依赖其提供的各种检查。

任务名称

支持的任务可能取决于后端，但这里是 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`	`自动模型`
`fill-mask`	`AutoModelForMaskedLM`
`问题回答`	`AutoModelForQuestionAnswering`
`text2text-generation`, `text2text-generation-with-past`	`AutoModelForSeq2SeqLM`
`文本分类`	`AutoModelForSequenceClassification`
`词元分类`	`AutoModelForTokenClassification`
`multiple-choice`	`AutoModelForMultipleChoice`
`image-classification`	`AutoModelForImageClassification`
`object-detection`	`AutoModelForObjectDetection`
`image-segmentation`	`AutoModelForImageSegmentation`
`masked-im`	`AutoModelForMaskedImageModeling`
`semantic-segmentation`	`AutoModelForSemanticSegmentation`
`自动语音识别`	`AutoModelForSpeechSeq2Seq`

TensorFlow

任务	自动类
`text-generation`, `text-generation-with-past`	`TFAutoModelForCausalLM`
`default`, `default-with-past`	`TFAutoModel`
`fill-mask`	`TFAutoModelForMaskedLM`
`问题回答`	`TFAutoModelForQuestionAnswering`
`text2text-generation`, `text2text-generation-with-past`	`TFAutoModelForSeq2SeqLM`
`文本分类`	`TFAutoModelForSequenceClassification`
`词元分类`	`TFAutoModelForTokenClassification`
`multiple-choice`	`TFAutoModelForMultipleChoice`
`semantic-segmentation`	`TFAutoModelForSemanticSegmentation`

参考

class optimum.exporters.TasksManager

< 源 >

( )

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

create_register

< 源 >

( 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

< 源 >

( 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[str, bool, NoneType] = None ) → str

参数

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

字符串

用于导出的框架。

确定用于导出的框架。

优先级顺序如下：

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

get_all_tasks

< 源 >

( ) → List

列表

所有可能的任务。

检索所有可能的任务。

get_exporter_config_constructor

< 源 >

( 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.Optional[typing.Dict[str, typing.Any]] = 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) — 模型对象的名称属性，仅用于异常消息。
exporter_config_kwargs (Optional[Dict[str, Any]], 默认为 None) — 在构建配置构造函数时将传递给导出器配置类的参数。
library_name (Optional[str], 默认为 None) — 模型的库名称。可以是“transformers”、“timm”、“diffusers”、“sentence_transformers”中的任意一个。

ExportConfigConstructor

请求后端的 `ExporterConfig` 构造函数。

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

get_model_class_for_task

< 源 >

( task: str framework: str = 'pt' model_type: typing.Optional[str] = None model_class_name: typing.Optional[str] = None library: str = 'transformers' )

参数

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

< 源 >

( 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[str, bool, 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 承载授权的令牌。如果为 True，将使用运行 huggingface-cli login 时生成的令牌（存储在 huggingface_hub.constants.HF_TOKEN_PATH 中）。
framework (Optional[str], 可选) — 用于导出的框架。有关未提供框架时的优先级，请参见 TasksManager.determine_framework。
torch_dtype (Optional[torch.dtype], 默认为 None) — 加载模型的数据类型。仅限 PyTorch 参数。
device (Optional[torch.device], 默认为 None) — 初始化模型的设备。仅限 PyTorch 参数。对于 PyTorch，默认为“cpu”。
library_name (Optional[str], 默认为 None) — 模型的库名称。可以是“transformers”、“timm”、“diffusers”、“sentence_transformers”中的任意一个。有关未提供库时的优先级，请参见 TasksManager.infer_library_from_model。
model_kwargs (Dict[str, Any], 可选) — 传递给模型 .from_pretrained() 方法的关键字参数。
library_name (Optional[str], 默认为 None) — 模型的库名称。可以是“transformers”、“timm”、“diffusers”、“sentence_transformers”中的任意一个。有关未提供库时的优先级，请参见 TasksManager.infer_library_from_model。

根据其名称和要启用的任务检索模型。

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

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

从模型类型检索 `任务 -> 导出器后端配置构造函数` 映射。

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[str, bool, NoneType] = None ) → str

参数

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

字符串

根据模型仓库、模型实例或模型类自动检测到的库名称。

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

infer_task_from_model

< source >

参数

model (Union[str, PreTrainedModel, TFPreTrainedModel, DiffusionPipeline, Type]) — 要推断任务的模型。这可以是HuggingFace Hub上的仓库名称、模型实例或模型类。
subfolder (str, optional, 默认为 "") — 如果模型文件位于Hugging Face Hub上的模型目录/仓库的子文件夹中，可以在此处指定子文件夹名称。
revision (Optional[str], 默认为 None) — revision 是要使用的特定模型版本。它可以是分支名称、标签名称或提交ID。
cache_dir (Optional[str], optional) — 如果不使用标准缓存，则为下载的预训练模型权重缓存目录的路径。
token (Optional[Union[bool,str]], 默认为 None) — 用作远程文件HTTP bearer授权的token。如果为True，将使用运行huggingface-cli login时生成的token（存储在huggingface_hub.constants.HF_TOKEN_PATH中）。
library_name (Optional[str], 默认为 None) — 模型的库名称。可以是“transformers”、“timm”、“diffusers”、“sentence_transformers”中的任何一个。如果未提供，请参见TasksManager.infer_library_from_model了解优先级。

字符串

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

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

standardize_model_attributes

< source >

( model: typing.Union[ForwardRef('PreTrainedModel'), ForwardRef('TFPreTrainedModel'), ForwardRef('DiffusionPipeline')] library_name: typing.Optional[str] = None )

参数

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

更新模型以进行导出。此函数适用于对不同库中的模型进行必要的更改，以遵循transformers风格。

< > 在 GitHub 上更新

Optimum

任务管理器

任务名称

PyTorch

TensorFlow

参考

class 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