保存具有自定义前向功能的 Bert 模型并将其置于 Huggingface 上

我创建了自己的 BertClassifier 模型，从预训练开始，然后添加由不同层组成的我自己的分类头。微调后，我想使用 model.save_pretrained() 保存模型，但是当我打印它并从预训练上传时，我看不到我的分类器头。 代码如下。如何保存模型上的所有结构并使其完全可访问 AutoModel.from_preatrained('folder_path')？ 。谢谢！

class BertClassifier(PreTrainedModel):
    """Bert Model for Classification Tasks."""
    config_class = AutoConfig
    def __init__(self,config, freeze_bert=True): #tuning only the head
        """
         @param    bert: a BertModel object
         @param    classifier: a torch.nn.Module classifier
         @param    freeze_bert (bool): Set `False` to fine-tune the BERT model
        """
        #super(BertClassifier, self).__init__()
        super().__init__(config)

        # Instantiate BERT model
        # Specify hidden size of BERT, hidden size of our classifier, and number of labels
        self.D_in = 1024 #hidden size of Bert
        self.H = 512
        self.D_out = 2
 
        # Instantiate the classifier head with some one-layer feed-forward classifier
        self.classifier = nn.Sequential(
            nn.Linear(self.D_in, 512),
            nn.Tanh(),
            nn.Linear(512, self.D_out),
            nn.Tanh()
        )
 


    def forward(self, input_ids, attention_mask):


         # Feed input to BERT
        outputs = self.bert(input_ids=input_ids,
                             attention_mask=attention_mask)
         
         # Extract the last hidden state of the token `[CLS]` for classification task
        last_hidden_state_cls = outputs[0][:, 0, :]
 
         # Feed input to classifier to compute logits
        logits = self.classifier(last_hidden_state_cls)
 
        return logits

configuration=AutoConfig.from_pretrained('Rostlab/prot_bert_bfd')
model = BertClassifier(config=configuration,freeze_bert=False)

微调后保存模型

model.save_pretrained('path')

加载微调模型

model = AutoModel.from_pretrained('path') 

加载后打印模型显示我有以下最后一层，并且缺少我的 2 个线性层：

 (output): BertOutput(
          (dense): Linear(in_features=4096, out_features=1024, bias=True)
          (LayerNorm): LayerNorm((1024,), eps=1e-12, elementwise_affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (adapters): ModuleDict()
          (adapter_fusion_layer): ModuleDict()
        )
      )
    )
  )
  (pooler): BertPooler(
    (dense): Linear(in_features=1024, out_features=1024, bias=True)
    (activation): Tanh()
  )
  (prefix_tuning): PrefixTuningPool(
    (prefix_tunings): ModuleDict()
  )
)

也许有问题config_class你里面的属性BertClassifier班级。根据文档，您需要创建一个继承 form 的附加配置类PretrainedConfig并初始化model_type属性与您的自定义模型的名称。

The BertClassifier's config_class必须与您的自定义配置类类型一致。 之后，您可以通过以下调用注册您的配置和模型：

AutoConfig.register('CustomModelName', CustomModelConfigClass)
AutoModel.register(CustomModelConfigClass, CustomModelClass)

并加载您的微调模型AutoModel.from_pretrained('YourCustomModelName')

基于您的代码的不完整示例可能如下所示：

class BertClassifierConfig(PretrainedConfig):
    model_type="BertClassifier"


class BertClassifier(PreTrainedModel):
    config_class = BertClassifierConfig
    # ...


configuration = BertClassifierConfig()
bert_classifier = BertClassifier(configuration)

# do your finetuning and save your custom model
bert_classifier.save_pretrained("CustomModels/BertClassifier")

# register your config and your model
AutoConfig.register("BertClassifier", BertClassifierConfig)
AutoModel.register(BertClassifierConfig, BertClassifier)

# load your model with AutoModel
bert_classifier_model = AutoModel.from_pretrained("CustomModels/BertClassifier")

打印模型输出应与此类似：

    (pooler): BertPooler(
      (dense): Linear(in_features=768, out_features=768, bias=True)
      (activation): Tanh()
    )
  )
  (classifier): Sequential(
    (0): Linear(in_features=1024, out_features=512, bias=True)
    (1): Tanh()
    (2): Linear(in_features=512, out_features=2, bias=True)
    (3): Tanh()
    (4): Linear(in_features=2, out_features=512, bias=True)
    (5): Tanh()
  )

希望这可以帮助。

https://huggingface.co/docs/transformers/custom_models#registering-a-model-with-custom-code-to-the-auto-classes https://huggingface.co/docs/transformers/custom_models#registering-a-model-with-custom-code-to-the-auto-classes