Including classification heads in BERT saves
Hi, I am using theTrainerwith BERT for a classification Task. I saved the model with this command:model = transformers.AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")
...
trainer.train()
_ = trainer.save_model("history/")Since BERT does not use a classification head, this code adds one, if I remember right. However, it appears that when I load this model after training withconfig = transformers.AutoConfig.from_pretrained("history")
model = transformers.AutoModelForSequenceClassification.from_config(config)
tokenizer = transformers.AutoTokenizer.from_pretrained("learning/checkpoint-31500")
pipe = transformers.TextClassificationPipeline(
	model=model,
	tokenizer=tokenizer,
	return_all_scores=True)It seems to create a new head again, not use the one it originally trained. This makes loading and using a trained model not possible.How can I save the head thatTrainertrained for bert within the model so I can load and use it again?Thanks in advance!
2023-04-06T19:45:37Z
Since HF models inherit from torch.nn.module, you can save your model with it’s custom classifier head with:torch.save(model, 'model.pt')and load it withtorch.load('model.pt'). Note that the loaded model can be input into themodelargument of thetrainerinstance.Now the second issue. How can you modifyTrainerso that it correctly saves your model checkpoints? This might not be the most elegant solution but I was able to get it by overriding the_savemethod ofTrainerlike this:class CustomTrainer(Trainer):
    def _save(self, output_dir, *args, **kwargs):
        output_dir = output_dir if output_dir is not None else self.args.output_dir
        os.makedirs(output_dir, exist_ok=True)
        torch.save(self.model, os.path.join(output_dir, 'model.pt'))Then creating your trainer instance with:# create Trainer instance
trainer = CustomTrainer(
    model=loaded_model,
    args=training_args,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    train_dataset=dataset['train'],
    eval_dataset=dataset['test'],
)andtrainerwill save a file calledmodel.ptin the checkpoint folder.Note, you will not be able to load the model using....from_pretrained()because it is not in the correct format but you will be able to load it with the torch load method above and supply that to yourtrainerinstance. Hope this helps!