将模型权重单独保存在pytorch中

Question

我正在使用 PyTorch 训练深度学习模型。我想知道我是否可以单独保存模型重量。例如：

class my_model(nn.Module):
def __init__(self):
    super(my_model, self).__init__()
    self.bert = transformers.AutoModel.from_pretrained(BERT_PATH)
    self.out = nn.Linear(768,1)
    
def forward(self, ids, mask, token_type):
    x = self.bert(ids, mask, token_type)[1]
    x = self.out(x)
    
    return x

我将 BERT 模型作为基础模型，并在顶部添加了一个线性层。我训练完这个模型后，能不能把BERT模型和这个线性层的权重分开保存？

Answer 1

您可以：

model = my_model()
# train ...
torch.save({'bert': model.bert.state_dict(), 'out': model.out.state_dict()}, 'checkpoint.pth')

Answer 2

与之前的答案不同，您可以创建两个分开的 class of nn.module。一张用于BERT模型，一张用于线性层：

class bert_model(nn.Module):
  def __init__(self):
  super(bert_model, self).__init__()
  self.bert = transformers.AutoModel.from_pretrained(BERT_PATH)

  def forward(self, ids, mask, token_type):
    x = self.bert(ids, mask, token_type)[1]

    return x

class linear_layer(nn.Module):
  def __init__(self):
  super(linear_layer, self).__init__()
  self.out = nn.Linear(768,1)

  def forward(self, x):
    x = self.out(x)

    return x

然后你可以分别保存模型的两个部分：

bert_model = bert_model()
linear_layer = linear_layer()
#train
torch.save(bert_model.state_dict(), PATH)
torch.save(linear_layer.state_dict(), PATH)

将模型权重单独保存在pytorch中

separately save the model weight in pytorch

python

machine-learning

deep-learning

pytorch

bert-language-model