由于内存问题,如何保存与预训练的bert模型的分类器层相关的参数?

How to save parameters just related to classifier layer of pretrained bert model due to the memory concerns?

我通过冻结除 classifier 层之外的所有层来微调预训练模型 here。我使用 pytorch 将权重文件保存为 .bin 格式。

现在不是加载 400mb 的预训练模型,有没有办法加载我重新训练的刚刚分类器层的参数?顺便说一句,我知道我必须加载原始的预训练模型,我只是不想加载整个微调模型。由于内存问题。

我可以从 state_dict 访问最后一层的参数,如下所示,但是我如何将它们保存在一个单独的文件中以便以后使用它们以减少内存使用量?

model = PosTaggingModel(num_pos_tag=num_pos_tag)
state_dict = torch.load("model.bin")
print("state dictionary:",state_dict)
with torch.no_grad():
    model.out_pos_tag.weight.copy_(state_dict['out_pos_tag.weight'])
    model.out_pos_tag.bias.copy_(state_dict['out_pos_tag.bias'])

这是模型class:

class PosTaggingModel(nn.Module):
    def __init__(self, num_pos_tag):
        super(PosTaggingModel, self).__init__()
        self.num_pos_tag = num_pos_tag
        self.model = AutoModel.from_pretrained("dbmdz/bert-base-turkish-cased")
        for name, param in self.model.named_parameters():
            if 'classifier' not in name: # classifier layer
                param.requires_grad = False
        self.bert_drop = nn.Dropout(0.3)
        self.out_pos_tag = nn.Linear(768, self.num_pos_tag)
        
    def forward(self, ids, mask, token_type_ids, target_pos_tag):
        o1, _ = self.model(ids, attention_mask = mask, token_type_ids = token_type_ids)
        
        bo_pos_tag = self.bert_drop(o1)
        pos_tag = self.out_pos_tag(bo_pos_tag)

        loss = loss_fn(pos_tag, target_pos_tag, mask, self.num_pos_tag)
        return pos_tag, loss

我不知道这是否可行,但我只是在寻找一种方法来保存和重用最后一层的参数,而不需要冻结层的参数。我在 documentation 中找不到它。 在此先感谢那些愿意提供帮助的人。

你可以这样做

import torch

# creating a dummy model
class Classifier(torch.nn.Module):
  def __init__(self):
    super(Classifier, self).__init__()
    self.first = torch.nn.Linear(10, 10)
    self.second = torch.nn.Linear(10, 20)
    self.last = torch.nn.Linear(20, 1)
  
  def forward(self, x):
    pass

# Creating its object
model = Classifier()

# Extracting the layer to save
to_save = model.last

# Saving the state dict of that layer
torch.save(to_save.state_dict(), './classifier.bin')

# Recreating the object of that model
model = Classifier()

# Updating the saved layer of model
model.last.load_state_dict(torch.load('./classifier.bin'))