将 model 权重单独保存在 pytorch 中

Question

I am using PyTorch to train a deep learning model. I wonder if it is possible for me to separately save the model weight. For example:

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

I have the BERT model as the base model and an additional linear layer on the top. After I train this model, can I save the weight for the BERT model and this linear layer separately?

Answer 1

You can:

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

Answer 2

Alternatively to the previous answer, You can create two separated class of nn.module. One for the BERT model and another one for the linear layer:

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

Then you can save the two part of the model separately with:

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