# Specify Loss for Trainer / TrainingArguments

**URL:** <https://discuss.huggingface.co/t/specify-loss-for-trainer-trainingarguments/10481>\
**Category:** 🤗Transformers\
**Created:** [October 3, 2021, 10:30pm UTC](https://discuss.huggingface.co/t/specify-loss-for-trainer-trainingarguments/10481 "2021-10-03T22:30:33Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![RylanSchaeffer](https://avatars.discourse-cdn.com/v4/letter/r/6f9a4e/32.png) [@RylanSchaeffer](https://discuss.huggingface.co/u/RylanSchaeffer)\
**Post date:** [October 3, 2021, 10:30pm UTC](https://discuss.huggingface.co/t/specify-loss-for-trainer-trainingarguments/10481/1 "2021-10-03T22:30:33Z")

</div>

I’d like to fine-tune for a regression task rather than a classification task. How do I change the default loss in either `TrainingArguments` or `Trainer`()?

---

<div class="post-metadata">

**Author:** ![nielsr](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/nielsr/32/39617_2.png) [@nielsr](https://discuss.huggingface.co/u/nielsr)\
**Post date:** [October 4, 2021, 8:34am UTC](https://discuss.huggingface.co/t/specify-loss-for-trainer-trainingarguments/10481/2 "2021-10-04T08:34:44Z")

</div>

You can overwrite the `compute_loss` method of the Trainer, like so:

```auto
from torch import nn
from transformers import Trainer

class RegressionTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.get("labels")
        outputs = model(**inputs)
        logits = outputs.get('logits')
        loss_fct = MSELoss()
        loss = loss_fct(logits.squeeze(), labels.squeeze())
        return (loss, outputs) if return_outputs else loss

```

However, [several models in the library](https://huggingface.co/transformers/master/main_classes/configuration.html?highlight=problem_type) have an attribute of their config called `problem_type`, which you can set to “regression”. In that case, you shouldn’t overwrite anything, and you can just use the default loss of the model.

---

<div class="post-metadata">

**Author:** ![RylanSchaeffer](https://avatars.discourse-cdn.com/v4/letter/r/6f9a4e/32.png) [@RylanSchaeffer](https://discuss.huggingface.co/u/RylanSchaeffer)\
**Post date:** [October 4, 2021, 3:36pm UTC](https://discuss.huggingface.co/t/specify-loss-for-trainer-trainingarguments/10481/3 "2021-10-04T15:36:10Z")

</div>

Thank you!

---

<div class="post-metadata">

**Author:** ![RylanSchaeffer](https://avatars.discourse-cdn.com/v4/letter/r/6f9a4e/32.png) [@RylanSchaeffer](https://discuss.huggingface.co/u/RylanSchaeffer)\
**Post date:** [October 5, 2021, 1:06am UTC](https://discuss.huggingface.co/t/specify-loss-for-trainer-trainingarguments/10481/4 "2021-10-05T01:06:47Z")

</div>

@nielsr , I tried this and the model’s predictions concentrate around a single predicted value. I’m almost exactly copying the fine-tuning tutorial. Any idea why the model seems to be failing to learn?

![](https://us1.discourse-cdn.com/hellohellohello/original/2X/3/3b6bcd99ec3412b120ab29919321359b141c57f9.png)

My code:

```auto
    training_args = TrainingArguments(
        output_dir=results_dir, # output directory
        num_train_epochs=50, # total number of training epochs
        per_device_train_batch_size=16, # batch size per device during training
        per_device_eval_batch_size=64, # batch size for evaluation
        warmup_steps=500, # number of warmup steps for learning rate scheduler
        weight_decay=0.01, # strength of weight decay
        logging_dir=results_dir, # directory for storing logs
        logging_steps=10,
        report_to='wandb',
        do_eval=True,
        evaluation_strategy="steps",
        eval_steps=10,
    )

    class RegressionTrainer(Trainer):

        def compute_loss(self,
                         model,
                         inputs,
                         return_outputs=False):
            labels = inputs.get("labels")
            outputs = model(**inputs)
            logits = outputs.get('logits')
            loss = torch.mean(torch.square(logits.squeeze() - labels.squeeze()))
            return (loss, outputs) if return_outputs else loss

    pytorch_model_save_path = os.path.join(results_dir, 'pytorch_model.bin')
    if os.path.isfile(pytorch_model_save_path): # If model was already fine-tuned
        # yes, pass the whole results dir; see https://github.com/huggingface/transformers/issues/1620
        model = DistilBertForSequenceClassification.from_pretrained(
            results_dir,
            num_labels=1)
    else: # If model needs to be fine-tuned
        # Set output dimension to 1 to perform regression
        model = DistilBertForSequenceClassification.from_pretrained(
            "distilbert-base-uncased",
            num_labels=1)

    trainer = RegressionTrainer(
        model=model, # the instantiated 🤗 Transformers model to be trained
        args=training_args, # training arguments, defined above
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        # compute_metrics=compute_eval_metrics,
    )

    if force_train or not os.path.isfile(pytorch_model_save_path):
        trainer.train()
        trainer.save_model(output_dir=results_dir)

    all_prediction_output = trainer.predict(all_dataset)

    all_predictions = all_prediction_output.predictions
    all_predictions = all_predictions.squeeze()
    all_labels = all_prediction_output.label_ids

```

---

<div class="post-metadata">

**Author:** ![RylanSchaeffer](https://avatars.discourse-cdn.com/v4/letter/r/6f9a4e/32.png) [@RylanSchaeffer](https://discuss.huggingface.co/u/RylanSchaeffer)\
**Post date:** [October 5, 2021, 1:09am UTC](https://discuss.huggingface.co/t/specify-loss-for-trainer-trainingarguments/10481/5 "2021-10-05T01:09:36Z")

</div>

This is weird because my validation loss is going down, but in a suspiciously smooth way:

![](https://us1.discourse-cdn.com/hellohellohello/original/2X/0/0c7dc0787328958be57639aa86e884cb7a94e287.png)

---

<div class="post-metadata">

**Author:** ![RylanSchaeffer](https://avatars.discourse-cdn.com/v4/letter/r/6f9a4e/32.png) [@RylanSchaeffer](https://discuss.huggingface.co/u/RylanSchaeffer)\
**Post date:** [October 5, 2021, 1:27am UTC](https://discuss.huggingface.co/t/specify-loss-for-trainer-trainingarguments/10481/6 "2021-10-05T01:27:36Z")

</div>

Turns out there was no error! Two things:

1. The learning rate was small and the validation loss was being evaluated very frequently, which explains why the validation loss was so smooth.

2. I needed to run 50 training epochs to see a real difference. That seems odd, but so be it