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 that Trainer trained for bert within the model so I can load and use it again?
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 with torch.load('model.pt'). Note that the loaded model can be input into the model argument of the trainer instance.
Now the second issue. How can you modify Trainer so 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 _save method of Trainer like this:
def _save(self, output_dir, *args, **kwargs):
output_dir = output_dir if output_dir is not None else self.args.output_dir
torch.save(self.model, os.path.join(output_dir, 'model.pt'))
and trainer will save a file called model.pt in 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 your trainer instance. Hope this helps!