separately save the model weight in pytorch-CodePudding

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?

CodePudding user response：

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)

CodePudding user response：

You can:

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