is there a way to save only the model with huggingface trainer?

user_2160809 0
Viewed 889

I want to keep multiple checkpoints during training to analyse them later but the Trainer also saves other files to resume training. Is there a way to only save the model to save space and writing time?

 15K  rng_state.pth
 906  trainer_state.json
 623  scheduler.pt
2,1G  optimizer.pt
2,5K  training_args.bin
1,1G  pytorch_model.bin
 900  config.json

I could just delete the optimizer after training but i'm also working with a disk with slower write speed so this is also a consideration

huggingface-transformers
1 Answers

Unfortunately, there is currently no way to disable the saving of single files. There are basically two ways to get your behavior:

  1. The "hacky" way would be to simply disable the line of code in the Trainer source code that stores the optimizer, which (if you train on your local machine) should be this one. However, whenever you update your transformers version, this could lead to ugly behavior, which is why I recommend the second one.
  2. Overwrite the _save_checkpoint() function in your own Trainer object. This way, you always guarantee that the correct files are saved, and don't have to interact with the library's code.

An outline for what this looks like:

from transformers import Trainer
class OwnTrainer(Trainer):
   # Don't forget to correctly set up __init__()
   # ...

   def _save_checkpoint(self, model, trial, metrics=None):
       # insert your own behavior here
user_36072030
Related