如何自动将自定义PyTorch模型转换为HuggingFace模型

问题场景 #

假设你已经创建了如下自定义PyTorch模型：

import torch
import torch.nn as nn
from transformers import RobertaModel

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        
        self.model_1_encoder = RobertaModel.from_pretrained('roberta-base')
        self.model_2_encoder = RobertaModel.from_pretrained('roberta-base')

        self.dropout = nn.Dropout(0.5)
        self.linear_model_1_out = nn.Linear(768, 512)
        self.linear_model_2_out = nn.Linear(768, 512)
    
        self.linear3 = nn.Linear(512, 512)
        self.linear4 = nn.Linear(512, 1)
        
    def forward(self, x):
        question_input_ids, attn_mask, xyz_input_ids, xyz_attention_mask = torch.tensor(x['input_ids']), torch.tensor(x['attention_mask']), torch.tensor(x['xyz_ids']), torch.tensor(x['xyz_mask'])
        model_1_outputs = self.model_1_encoder(question_input_ids, attention_mask=attn_mask)
        model_2_outputs = self.model_2_encoder(xyz_input_ids, attention_mask = xyz_attention_mask)
        model_1_outputs = self.dropout(model_1_outputs[0])
        model_2_outputs = self.dropout(model_2_outputs[0])

        model_2_outputs = torch.sum(model_2_outputs, dim=0) 
        model_2_outputs = torch.unsqueeze(model_2_outputs, 0) 
        model_2_outputs = model_2_outputs.repeat(model_1_outputs.shape[0], 1, 1) 

        model_1_outputs = self.linear_model_1_out(model_1_outputs)
        model_2_outputs = self.linear_model_2_out(model_2_outputs)

        model_1_and_xyz_sum = model_2_outputs + model_1_outputs 
        outputs = self.dropout(model_1_and_xyz_sum)
        outputs = outputs[:,0,:].view(-1,512) 
        outputs = torch.relu(self.linear3(outputs))
        outputs = self.dropout(outputs)
        outputs = torch.relu(self.linear4(outputs))
        return outputs

model = Model()

常规自定义模型接入HuggingFace的流程需要手动编写配置类、逐层对齐模型结构、完成多项注册配置，步骤繁琐。实际上不需要走完整的自定义模型注册流程，只要让自定义模型继承PreTrainedModel基类，就能直接复用HuggingFace生态的Trainer、DataCollator、from_pretrained/save_pretrained等所有内置功能，完全匹配极简调用的需求。

具体实现步骤 #

• 第一步：定义和模型匹配的极简配置类，继承PretrainedConfig，只需要传入模型用到的超参数即可，不需要额外编写复杂逻辑
• 第二步：把原有模型类修改为继承PreTrainedModel，在__init__方法首行调用父类初始化传入config，原有模型结构、前向传播逻辑几乎不需要改动
• 第三步：添加一个一键加载的类方法，直接加载原有PyTorch模型的权重，实例化后即可无缝接入HuggingFace全生态工具

可直接运行的实现代码 #

from transformers import PreTrainedModel, PretrainedConfig

# 定义极简配置类
class CustomModelConfig(PretrainedConfig):
    model_type = "custom_dual_roberta"
    def __init__(self, dropout_rate=0.5, hidden_size=768, proj_size=512, num_labels=1, **kwargs):
        super().__init__(**kwargs)
        self.dropout_rate = dropout_rate
        self.hidden_size = hidden_size
        self.proj_size = proj_size
        self.num_labels = num_labels

# 包装原有模型
class CustomHuggingfaceModel(PreTrainedModel):
    config_class = CustomModelConfig
    def __init__(self, config):
        super().__init__(config)
        # 原有模型初始化逻辑无需改动，仅把硬编码参数替换为config读取即可
        self.model_1_encoder = RobertaModel.from_pretrained('roberta-base')
        self.model_2_encoder = RobertaModel.from_pretrained('roberta-base')
        self.dropout = nn.Dropout(config.dropout_rate)
        self.linear_model_1_out = nn.Linear(config.hidden_size, config.proj_size)
        self.linear_model_2_out = nn.Linear(config.hidden_size, config.proj_size)
        self.linear3 = nn.Linear(config.proj_size, config.proj_size)
        self.linear4 = nn.Linear(config.proj_size, config.num_labels)
        # 调用内置方法完成权重初始化
        self.post_init()

    # 原有前向传播逻辑无需改动，仅增加labels可选参数适配Trainer自动算损失
    def forward(self, x, labels=None):
        question_input_ids, attn_mask, xyz_input_ids, xyz_attention_mask = x['input_ids'], x['attention_mask'], x['xyz_ids'], x['xyz_mask']
        model_1_outputs = self.model_1_encoder(question_input_ids, attention_mask=attn_mask)
        model_2_outputs = self.model_2_encoder(xyz_input_ids, attention_mask = xyz_attention_mask)
        model_1_outputs = self.dropout(model_1_outputs[0])
        model_2_outputs = self.dropout(model_2_outputs[0])

        model_2_outputs = torch.sum(model_2_outputs, dim=0) 
        model_2_outputs = torch.unsqueeze(model_2_outputs, 0) 
        model_2_outputs = model_2_outputs.repeat(model_1_outputs.shape[0], 1, 1) 

        model_1_outputs = self.linear_model_1_out(model_1_outputs)
        model_2_outputs = self.linear_model_2_out(model_2_outputs)

        model_1_and_xyz_sum = model_2_outputs + model_1_outputs 
        outputs = self.dropout(model_1_and_xyz_sum)
        outputs = outputs[:,0,:].view(-1, self.config.proj_size) 
        outputs = torch.relu(self.linear3(outputs))
        outputs = self.dropout(outputs)
        logits = torch.relu(self.linear4(outputs))

        loss = None
        if labels is not None:
            loss_fct = nn.MSELoss()
            loss = loss_fct(logits.view(-1), labels.view(-1))
        return {"loss": loss, "logits": logits}

    # 实现期望的一键加载逻辑
    @classmethod
    def load_model(cls, existing_pytorch_model=None, config=None):
        if config is None:
            config = CustomModelConfig()
        model = cls(config)
        if existing_pytorch_model is not None:
            # 直接加载原有PyTorch模型的权重，无需重新训练
            model.load_state_dict(existing_pytorch_model.state_dict())
        return model

调用方式 #

完全匹配期望的伪代码逻辑：

# 原有训练好的PyTorch模型
my_pytorch_model = Model()
# 一键转换为HuggingFace兼容模型
my_huggingface_model = CustomHuggingfaceModel.load_model(my_pytorch_model)

# 后续可直接使用HuggingFace所有内置功能
# 比如保存/加载模型
my_huggingface_model.save_pretrained("./my_custom_model")
loaded_model = CustomHuggingfaceModel.from_pretrained("./my_custom_model")

注意：如果前向传播方法原来接收字典输入，使用DataCollator时只要保证返回的字典字段和前向传播里读取的key一致即可，不需要额外调整数据处理逻辑。

内容的提问来源于stack exchange，提问作者Penguin