python - 由于内存问题，如何保存仅与预训练bert模型的分类器层相关的参数？

Question

我微调了预训练模型 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 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 中找不到它.
在此先感谢那些愿意提供帮助的人。

Answer 1

你可以这样做

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'))