How to use a custom loss with hugging face

Recently I helped a colleague at work to fix an issue to setup a custom loss when using hugging face. We couldn’t find much information about it, so we had to figure it out ourselves. Here I want to share it so others can find it next time.

Let's start by setting up a dummy model, dataset, and hugging face trainer (forget about optimality, or if it makes sense at all, we just want to make sure it works end to end without failing).

Ok, so we are using IMDB dataset to classify reviews as good or bad. This a classification problem and hugging face correctly setup the loss function and everything else for us… But! Let's go the hard way, now we are going to use our own loss function, how?

Subclass Trainer and implement the method compute_loss:

What if my custom trainer is not receiving the labels?

This happened to us. You might have “remove_unused_columns=True”, meaning the trainer is dropping the inputs that are not expected on the forward function, so either set “remove_unused_columns=False” or add the labels arguments to your module:

How that works is by using “inspect.signature”, which gives you the ability to read the parameters of the given function: https://github.com/huggingface/transformers/blob/f1660d7e23d4432513fe060bde4f9b7b29f05204/src/transformers/trainer.py#L755