有没有办法只保存带有 Huggingface 训练器的模型？

Question

我想在训练期间保留多个检查点以便稍后分析它们，但Trainer也保存其他文件以恢复训练。有没有办法只保存模型来节省空间和编写时间？

 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

我可以在训练后删除优化器，但我也在使用写入速度较慢的磁盘，所以这也是一个考虑因素

Answer 1

不幸的是，目前没有办法禁用单个文件的保存。基本上有两种方法可以获取您的行为：

1. Trainer“hacky”方法是简单地禁用源代码中存储优化器的代码行，（如果您在本地计算机上训练）应该是这一行。然而，每当你更新你的变形金刚版本时，这都可能导致丑陋的行为，这就是我推荐第二个的原因。
2. 覆盖_save_checkpoint()您自己的Trainer对象中的函数。这样，您始终可以保证保存正确的文件，并且不必与库的代码进行交互。

大致如下：

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