How to save and load fine-tune model

🤗Transformers
1

Hi, everyone~ I have defined my model via huggingface, but I don’t know how to save and load the model, hopefully someone can help me out, thanks!

class MyModel(nn.Module):
  def __init__(self, num_classes):
    super(MyModel, self).__init__()
    self.bert = BertModel.from_pretrained('hfl/chinese-roberta-wwm-ext', return_dict=True).to(device)
    self.fc = nn.Linear(768, num_classes, bias=False)
  
  def forward(self, x_input_ids, x_type_ids, attn_mask):
    outputs = self.bert(x_input_ids, token_type_ids=x_type_ids, attention_mask=attn_mask)
    pred = self.fc(outputs.pooler_output)
    return pred

model = MyModel(num_classes).to(device)
# save 

# load
1 Like