Change saving metric in Trainer

I am training a model while performing hyperparameter optimization with optuna, therefore I can’t afford saving a lot of checkpoints in each trial because my machine would run out of disk. For this reason in each trial I am setting save_total_limit=2. This means that the two checkpoints with less validation loss will be saved for each trial. How can I do if I want to save the two checkpoints which have best f1 instead of the least validation loss? Is there any parameter or I have to inherit Trainer and make my own code?

I don’t know if there is a way to do that but there is load_best_model_at_end=True. But I’m not quite sure if it is something you want. Maybe see here: Save only best model in Trainer - #5 by sgugger

Hi,

You can use the argument metric_for_best_model in the Trainer to define what metric to use when choosing the best checkpoint. Do not forget to set the greater_is_better parameter, too, depending on your metric. Since you are using the F1 score, it should be True. Hope this helps!

1 Like