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
2

If you make your model a subclass of PreTrainedModel, then you can use our methods save_pretrained and from_pretrained. Otherwise it’s regular PyTorch code to save and load (using torch.save and torch.load).

2 Likes
3

What if the pre-trained model is saved by using torch.save(model.state_dict()). How can I use that model like the BertTokenizer for creating tokens and also embeddings?

4

Hi,
Is there a difference between loading a model via torch.load and using from_pretrained in terms of downstream tasks?
Does either method have an advantage over the other for fine-tuning or inference?

2 Likes
5

If i do torch.save it will save only the model file but it won’t save the config.json file . How to achieve that using torch.save method

6 Likes