Transformers 文档

构建自定义模型

Hugging Face's logo
加入 Hugging Face 社区

并获取增强型文档体验

开始使用

构建自定义模型

🤗 Transformers 库旨在易于扩展。每个模型都在仓库的给定子文件夹中完全编码,没有抽象,因此您可以轻松地复制模型文件并根据您的需要进行调整。

如果您正在编写全新的模型,从头开始可能更容易。在本教程中,我们将向您展示如何编写自定义模型及其配置,以便它可以在 Transformers 中使用,以及如何与社区共享它(以及它所依赖的代码),以便任何人都可以使用它,即使它不在 🤗 Transformers 库中。我们将看到如何构建在 transformers 之上,并使用您的钩子和自定义代码扩展框架。

我们将通过将 timm 库 的 ResNet 类封装到 PreTrainedModel 中,在 ResNet 模型上说明所有这些内容。

编写自定义配置

在我们深入模型之前,让我们先编写它的配置。模型的配置是一个对象,它将包含构建模型所需的所有必要信息。正如我们将在下一节中看到的那样,模型只能接收一个 config 来初始化,因此我们需要这个对象尽可能完整。

transformers 库本身中的模型通常遵循这样的约定:它们在 __init__ 方法中接受一个 config 对象,然后将整个 config 传递给模型中的子层,而不是将 config 对象分解成多个参数,这些参数都单独传递给子层。以这种方式编写您的模型会导致代码更简单,并为任何超参数提供清晰的“真相来源”,也使它更容易重用 transformers 中其他模型的代码。

在我们的示例中,我们将采用 ResNet 类中我们可能希望调整的几个参数。不同的配置将给我们不同的 ResNet 类型。然后我们只存储这些参数,并在检查其中一些参数的有效性后。

from transformers import PretrainedConfig
from typing import List


class ResnetConfig(PretrainedConfig):
    model_type = "resnet"

    def __init__(
        self,
        block_type="bottleneck",
        layers: List[int] = [3, 4, 6, 3],
        num_classes: int = 1000,
        input_channels: int = 3,
        cardinality: int = 1,
        base_width: int = 64,
        stem_width: int = 64,
        stem_type: str = "",
        avg_down: bool = False,
        **kwargs,
    ):
        if block_type not in ["basic", "bottleneck"]:
            raise ValueError(f"`block_type` must be 'basic' or bottleneck', got {block_type}.")
        if stem_type not in ["", "deep", "deep-tiered"]:
            raise ValueError(f"`stem_type` must be '', 'deep' or 'deep-tiered', got {stem_type}.")

        self.block_type = block_type
        self.layers = layers
        self.num_classes = num_classes
        self.input_channels = input_channels
        self.cardinality = cardinality
        self.base_width = base_width
        self.stem_width = stem_width
        self.stem_type = stem_type
        self.avg_down = avg_down
        super().__init__(**kwargs)

编写自己的配置时,要记住三个重要事项:

  • 您必须继承自 PretrainedConfig
  • 您的 PretrainedConfig__init__ 必须接受任何 kwargs,
  • 这些 kwargs 需要传递给超类 __init__

继承是为了确保您从 🤗 Transformers 库中获得所有功能,而另外两个约束来自 PretrainedConfig 拥有比您正在设置的字段更多的字段。当使用 from_pretrained 方法重新加载配置时,这些字段需要被您的配置接受,然后发送到超类。

为您的配置定义 model_type(此处为 model_type="resnet")不是必需的,除非您想将您的模型注册到自动类中(见最后一节)。

完成此操作后,您可以像处理库中的任何其他模型配置一样轻松地创建和保存您的配置。以下是我们如何创建一个 resnet50d 配置并保存它:

resnet50d_config = ResnetConfig(block_type="bottleneck", stem_width=32, stem_type="deep", avg_down=True)
resnet50d_config.save_pretrained("custom-resnet")

这将在 custom-resnet 文件夹中保存一个名为 config.json 的文件。然后,您可以使用 from_pretrained 方法重新加载您的配置

resnet50d_config = ResnetConfig.from_pretrained("custom-resnet")

您还可以使用 PretrainedConfig 类的任何其他方法,例如 push_to_hub(),直接将您的配置上传到 Hub。

编写自定义模型

现在我们有了 ResNet 配置,我们可以继续编写模型。我们实际上会编写两个:一个从一批图像中提取隐藏特征(例如 BertModel),另一个适合图像分类(例如 BertForSequenceClassification)。

正如我们之前提到的,我们将只编写模型的松散包装器,以便在本例中保持简单。在编写此类之前,我们唯一需要做的事情是块类型与实际块类之间的映射。然后通过将所有内容传递给 ResNet 类来从配置中定义模型

from transformers import PreTrainedModel
from timm.models.resnet import BasicBlock, Bottleneck, ResNet
from .configuration_resnet import ResnetConfig


BLOCK_MAPPING = {"basic": BasicBlock, "bottleneck": Bottleneck}


class ResnetModel(PreTrainedModel):
    config_class = ResnetConfig

    def __init__(self, config):
        super().__init__(config)
        block_layer = BLOCK_MAPPING[config.block_type]
        self.model = ResNet(
            block_layer,
            config.layers,
            num_classes=config.num_classes,
            in_chans=config.input_channels,
            cardinality=config.cardinality,
            base_width=config.base_width,
            stem_width=config.stem_width,
            stem_type=config.stem_type,
            avg_down=config.avg_down,
        )

    def forward(self, tensor):
        return self.model.forward_features(tensor)

对于将对图像进行分类的模型,我们只需更改 forward 方法

import torch


class ResnetModelForImageClassification(PreTrainedModel):
    config_class = ResnetConfig

    def __init__(self, config):
        super().__init__(config)
        block_layer = BLOCK_MAPPING[config.block_type]
        self.model = ResNet(
            block_layer,
            config.layers,
            num_classes=config.num_classes,
            in_chans=config.input_channels,
            cardinality=config.cardinality,
            base_width=config.base_width,
            stem_width=config.stem_width,
            stem_type=config.stem_type,
            avg_down=config.avg_down,
        )

    def forward(self, tensor, labels=None):
        logits = self.model(tensor)
        if labels is not None:
            loss = torch.nn.functional.cross_entropy(logits, labels)
            return {"loss": loss, "logits": logits}
        return {"logits": logits}

在这两种情况下,请注意我们是如何继承自 PreTrainedModel 并使用 config 调用超类初始化的(有点像您编写常规 torch.nn.Module 时)。设置 config_class 的行不是必需的,除非您想将您的模型注册到自动类中(见最后一节)。

如果您的模型与库中的模型非常相似,您可以重用与该模型相同的配置。

您可以让您的模型返回任何您想要的内容,但像我们对 ResnetModelForImageClassification 所做的那样返回一个字典,包括传递标签时的损失,将使您的模型可以直接在 Trainer 类中使用。使用其他输出格式也可以,只要您计划使用自己的训练循环或其他库进行训练。

现在我们有了模型类,让我们创建一个

resnet50d = ResnetModelForImageClassification(resnet50d_config)

同样,您可以使用 PreTrainedModel 的任何方法,例如 save_pretrained()push_to_hub()。我们将在下一节中使用第二个方法,看看如何将模型权重与模型代码一起推送。但首先,让我们在模型中加载一些预训练的权重。

在您自己的用例中,您可能会在自己的数据上训练您的自定义模型。为了在本教程中快速完成,我们将使用 resnet50d 的预训练版本。由于我们的模型只是它的包装器,因此将这些权重转移过来将很容易

import timm

pretrained_model = timm.create_model("resnet50d", pretrained=True)
resnet50d.model.load_state_dict(pretrained_model.state_dict())

现在让我们看看如何确保当我们执行 save_pretrained()push_to_hub() 时,模型代码也被保存。

将具有自定义代码的模型注册到自动类中

如果您正在编写扩展 🤗 Transformers 的库,您可能希望扩展自动类以包含您自己的模型。这与将代码推送到 Hub 不同,因为用户需要导入您的库才能获得自定义模型(与从 Hub 自动下载模型代码相反)。

只要您的配置具有与现有模型类型不同的 model_type 属性,并且您的模型类具有正确的 config_class 属性,您就可以像这样将它们添加到自动类中

from transformers import AutoConfig, AutoModel, AutoModelForImageClassification

AutoConfig.register("resnet", ResnetConfig)
AutoModel.register(ResnetConfig, ResnetModel)
AutoModelForImageClassification.register(ResnetConfig, ResnetModelForImageClassification)

请注意,将您的自定义配置注册到 AutoConfig 时使用的第一个参数需要与您的自定义配置的 model_type 匹配,而将您的自定义模型注册到任何自动模型类时使用的第一个参数需要与这些模型的 config_class 匹配。

将代码发送到 Hub

此 API 处于实验阶段,在接下来的版本中可能会有一些细微的破坏性更改。

首先,确保你的模型在 .py 文件中完全定义。它可以依赖于对其他文件的相对导入,只要所有文件都在同一个目录中(目前我们不支持子模块)。例如,我们在当前工作目录下的一个名为 resnet_model 的文件夹中定义一个 modeling_resnet.py 文件和一个 configuration_resnet.py 文件。配置文件包含 ResnetConfig 的代码,而模型文件包含 ResnetModelResnetModelForImageClassification 的代码。

.
└── resnet_model
    ├── __init__.py
    ├── configuration_resnet.py
    └── modeling_resnet.py

__init__.py 可以为空,它只是用来让 Python 识别 resnet_model 可以用作模块。

如果从库中复制模型文件,你需要将文件顶部的所有相对导入替换为从 transformers 包中导入。

注意,你可以重用(或子类化)现有的配置/模型。

要与社区分享你的模型,请按照以下步骤操作:首先从新创建的文件中导入 ResNet 模型和配置。

from resnet_model.configuration_resnet import ResnetConfig
from resnet_model.modeling_resnet import ResnetModel, ResnetModelForImageClassification

然后,你必须告诉库,在使用 save_pretrained 方法时,你想复制这些对象的代码文件,并使用给定的 Auto 类(特别是对于模型)正确注册它们,只需运行

ResnetConfig.register_for_auto_class()
ResnetModel.register_for_auto_class("AutoModel")
ResnetModelForImageClassification.register_for_auto_class("AutoModelForImageClassification")

注意,对于配置不需要指定 Auto 类(它们只有一个 Auto 类,AutoConfig),但对于模型则不同。你的自定义模型可能适用于许多不同的任务,因此你需要指定哪个 Auto 类适合你的模型。

如果你想复制代码文件,请使用 register_for_auto_class()。如果你想从另一个仓库的 Hub 中使用代码,则不需要调用它。在有多个 Auto 类的情况下,你可以使用以下结构直接修改 config.json 文件

"auto_map": {     
	"AutoConfig": "<your-repo-name>--<config-name>",     
	"AutoModel": "<your-repo-name>--<config-name>",
	"AutoModelFor<Task>": "<your-repo-name>--<config-name>",    
},

接下来,让我们像以前一样创建配置和模型。

resnet50d_config = ResnetConfig(block_type="bottleneck", stem_width=32, stem_type="deep", avg_down=True)
resnet50d = ResnetModelForImageClassification(resnet50d_config)

pretrained_model = timm.create_model("resnet50d", pretrained=True)
resnet50d.model.load_state_dict(pretrained_model.state_dict())

现在,要将模型发送到 Hub,请确保你已登录。在终端中运行

huggingface-cli login

或在笔记本中运行

from huggingface_hub import notebook_login

notebook_login()

然后,你可以将模型推送到你自己的命名空间(或你所在的组织)中,如下所示

resnet50d.push_to_hub("custom-resnet50d")

除了模型权重和 JSON 格式的配置之外,这还会将 custom-resnet50d 文件夹中的模型和配置 .py 文件复制到 Hub。你可以在 模型库 中查看结果。

有关推送到 Hub 方法的更多信息,请查看 分享教程

使用带有自定义代码的模型

你可以使用任何带有自定义代码文件的仓库中的配置、模型或分词器,以及 Auto 类和 from_pretrained 方法。上传到 Hub 的所有文件和代码都会被扫描以检查恶意软件(有关更多信息,请参阅 Hub 安全 文档),但你仍应审查模型代码和作者,以避免在你的机器上执行恶意代码。设置 trust_remote_code=True 以使用带有自定义代码的模型。

from transformers import AutoModelForImageClassification

model = AutoModelForImageClassification.from_pretrained("sgugger/custom-resnet50d", trust_remote_code=True)

还强烈建议传递一个提交哈希作为 revision,以确保模型的作者没有使用一些恶意的新行更新代码(除非你完全信任模型的作者)。

commit_hash = "ed94a7c6247d8aedce4647f00f20de6875b5b292"
model = AutoModelForImageClassification.from_pretrained(
    "sgugger/custom-resnet50d", trust_remote_code=True, revision=commit_hash
)

注意,在浏览 Hub 上模型仓库的提交历史时,有一个按钮可以轻松复制任何提交的提交哈希。

< > 更新 在 GitHub 上