如何将 PyTorch nn.Module 转换为 HuggingFace PreTrainedModel 对象？

Question

给定 Pytorch 中的一个简单神经网络，例如：

import torch.nn as nn

net = nn.Sequential(
      nn.Linear(3, 4),
      nn.Sigmoid(),
      nn.Linear(4, 1),
      nn.Sigmoid()
      ).to(device)

如何将其转换为 Huggingface PreTrainedModel对象？

目标是将 Pytorchnn.Module对象转换nn.Sequential为 HuggingfacePreTrainedModel对象，然后运行如下所示的代码：

import torch.nn as nn
from transformers.modeling_utils import PreTrainedModel


net = nn.Sequential(
      nn.Linear(3, 4),
      nn.Sigmoid(),
      nn.Linear(4, 1),
      nn.Sigmoid()
      ).to(device)

# Do something to convert the Pytorch nn.Module to the PreTrainedModel object.
shiny_model = do_some_magic(net, some_args, some_kwargs)

# Save the shiny model that is a `PreTrainedModel` object.
shiny_model.save_pretrained("shiny-model")

PreTrainedModel.from_pretrained("shiny-model")

似乎要将任何本机 Pytorch 模型构建/转换为 Huggingface 模型，需要一些配置https://huggingface.co/docs/transformers/main_classes/configuration

有很多“从头开始”训练模型的方法，例如

• [使用 BertLMHeadModel，不是从头开始] https://www.kaggle.com/code/mojammel/train-model-from-scratch-with-huggingface/notebook（这也是从 bert 开始微调，不是从头开始）

• [不是真正从头开始，使用 roberta 作为模板] https://huggingface.co/blog/how-to-train（这是 roberta 的微调，不是真正从头开始训练）

• [有点使用一些配置模板] https://www.thepythoncode.com/article/pretraining-bert-huggingface-transformers-in-python（这有点从头开始，但使用 BERT 的模板来生成配置，如果我们想要改变模型的工作方式，配置应该是什么样子？）

• [有点定义了一个模板，但使用了 RobertaForMaskedLM] https://skimai.com/roberta-language-model-for-spanish/（这看起来有点定义了一个模板，但将其限制为 RobertaForMaskedLM 模板）

部分问题：

• 如果我们有一个像上面的代码片段一样简单得多的 Pytorch 模型，如何在 Huggingface 中从头开始创建预训练模型？

• 如何创建 Huggingface 所需的预训练模型配置，以便从本机 Pytorch nn.Module 进行转换？

Answer 1

您将需要定义自定义配置和自定义模型类。model_type在这些类中定义属性非常重要config_class：

import torch.nn as nn
from transformers import PreTrainedModel, PretrainedConfig
from transformers import AutoModel, AutoConfig

class MyConfig(PretrainedConfig):
    model_type = 'mymodel'
    def __init__(self, important_param=42, **kwargs):
        super().__init__(**kwargs)
        self.important_param = important_param

class MyModel(PreTrainedModel):
    config_class = MyConfig
    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.model = nn.Sequential(
                          nn.Linear(3, self.config.important_param),
                          nn.Sigmoid(),
                          nn.Linear(self.config.important_param, 1),
                          nn.Sigmoid()
                          )
    def forward(self, input):
        return self.model(input) 

现在您可以创建（显然是训练一个新模型）、保存然后在本地加载模型

config = MyConfig(4)
model = MyModel(config)
model.save_pretrained('./my_model_dir')

new_model = MyModel.from_pretrained('./my_model_dir')
new_model

如果您想使用AutoModel，您必须注册您的课程：

AutoConfig.register("mymodel", MyConfig)
AutoModel.register(MyConfig, MyModel)

new_model = AutoModel.from_pretrained('./my_model_dir')
new_model

Answer 2

一种方法是将模型放入继承自 PreTrainedModel 的类中，例如，它可以是预训练的 resnet34、timm 模型或您的“net”模型。我建议查看文档以获取有关配置的更多详细信息，我将使用链接中的示例。https://huggingface.co/docs/transformers/custom_models#sharing-custom-models

配置（注意：您可以为示例版本添加不同的配置，稍后可以访问 config.json。）

from transformers import PretrainedConfig
from typing import List

class ModelConfig(PretrainedConfig):
    model_type = "mymodel"
    def __init__(
        self,
        version = 1,
        layers: List[int] = [3, 4, 6, 3],
        num_classes: int = 1000,
        input_channels: int = 3,
        stem_type: str = "",
        **kwargs,
    ):
        if stem_type not in ["", "deep", "deep-tiered"]:
            raise ValueError(f"`stem_type` must be '', 'deep' or 'deep-tiered', got {block}.")

        self.version = version
        self.layers = layers
        self.num_classes = num_classes
        self.input_channels = input_channels
        self.stem_type = stem_type
        super().__init__(**kwargs)

正如我所说，你的网络模型可能是 resnet34。

from transformers import PreTrainedModel
from torch import nn
net = nn.Sequential(
      nn.Linear(3, 4),
      nn.Sigmoid(),
      nn.Linear(4, 1),
      nn.Sigmoid()
      ).to('cuda')
      
class MyModel(PreTrainedModel):
    config_class = ModelConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = net
        
    def forward(self, tensor):
        return self.model(tensor)

测试模型

config = ModelConfig()
model = MyModel(config)
dummy_input = torch.randn(1, 3).to('cuda')
with torch.no_grad():
    output = model(dummy_input)
print(output.shape)

推送到hugginface hub（注：需要使用token登录，可以多次推送更新模型）

model.push_to_hub("mymodel-test")

下载模型（注意：您正在使用MyModel类，如果您想创建像..bert.modeling_bert.BertModel这样的模型，我认为您需要使用lib结构。）

my_model = MyModel.from_pretrained("User/mymodel-test")