在HuggingFace里创建自己的模型

本文介绍了如何在Transformers库中扩展模型,通过实例化ResNet模型并创建自定义配置,以及如何在社区中共享这些代码。作者展示了如何编写配置文件,创建ResNetModel和ResnetModelForImageClassification,并加载预训练权重。
摘要由CSDN通过智能技术生成

本文为搬运huggingface文档:共享自定义模型

共享自定义模型

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

如果你要编写全新的模型,从头开始可能更容易。在本教程中,我们将向你展示如何编写自定义模型及其配置,以便可以在 Transformers 中使用它;以及如何与社区共享它(及其依赖的代码),以便任何人都可以使用,即使它不在 🤗 Transformers 库中。

我们将以 ResNet 模型为例,通过将 timm 库 的 ResNet 类封装到 PreTrainedModel 中来进行说明。

编写自定义配置

在深入研究模型之前,让我们首先编写其配置。模型的配置是一个对象,其中包含构建模型所需的所有信息。我们将在下一节中看到,模型只能接受一个 config 来进行初始化,因此我们很需要使该对象尽可能完整。

我们将采用一些我们可能想要调整的 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 方法重新加载配置:

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

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

编写自定义模型

有了 ResNet 配置后,就可以继续编写模型了。实际上,我们将编写两个模型:一个模型用于从一批图像中提取隐藏特征(类似于 BertModel),另一个模型适用于图像分类(类似于 BertForSequenceClassification)。

正如之前提到的,我们只会编写一个松散的模型包装,以使示例保持简洁。在编写此类之前,只需要建立起块类型(block types)与实际块类(block classes)之间的映射。然后,通过将所有内容传递给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)

对用于进行图像分类的模型,我们只需更改前向方法:

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.cross_entropy(logits, labels)
            return {"loss": loss, "logits": logits}
        return {"logits": logits}

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

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

你可以让模型返回任何你想要的内容,但是像我们为 ResnetModelForImageClassification 做的那样返回一个字典,并在传递标签时包含loss,可以使你的模型能够在 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())

(后面省略)

  • 2
    点赞
  • 3
    收藏
    觉得还不错? 一键收藏
  • 2
    评论
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值