起因
在训练模型时,要每隔一定步数 要验证一次,如果验证指标更好了,则要保存对应的checkpoints。但在实际模型训练过程中,我们不仅需要保存对应的checkpoint,还要删除最开始不用的那些checkpoints文件。
解决办法
之前在训练时使用的是transformers库里面的trainer模块,可以在定义TrainingArguments时设置save_total_limit 的数量来限制 checkpoints 的数量。
但是很多时候由于自己要魔改训练过程,需要自己手写训练过程,没办法使用transformer的trainer模块,所以我自己写了一个小函数,在每次验证结束要保存模型前调用即可。
def keep_checkpoints_num(ckpt_dir, ckpt_filename_keyword='checkpoint', ckpt_nums=3):
"""
维护ckpt_dir文件夹里checkpoints的数量
:param ckpt_dir: 存储checkpoints的文件夹
:param ckpt_filename_keyword: 单个checkpoint文件夹名字所包含的keyword,用来挑选出ckpt的文件夹名称
:param ckpt_nums: 存储ckpt的数量
:return:
"""
import shutil
from pathlib import Path
def get_ckpt_dirlist(ckpt_dir, ckpt_filename_keyword):
ckpt_dirlist = []
# 找到所有ckpt的文件夹名
for ckpt in Path(ckpt_dir).iterdir():
if ckpt_filename_keyword in ckpt.name:
ckpt_dirlist.append(ckpt)
return ckpt_dirlist
ckpt_dirlist = get_ckpt_dirlist(ckpt_dir, ckpt_filename_keyword)
# 按照生成时间进行排序,最近生成的排在最开始
ckpt_dirlist = sorted(ckpt_dirlist, key = lambda ckpt: ckpt.stat().st_mtime, reverse=True)
# 移除不满足条件的ckpt
remove_ckpt_dirlist = ckpt_dirlist[ckpt_nums - 1:]
for remove_ckpt in remove_ckpt_dirlist:
shutil.rmtree(remove_ckpt)
print(f"{remove_ckpt} is removed!")
具体使用:
如下图所示,我每隔200步验证一次模型,并保存。随着不断迭代会逐渐删除多余的ckpt,保存下来的是top ckpt_num 个模型。
其中ckpt_filename_keyword 可根据自定义的ckpt文件名设置,主要是筛选。
# ckpt_dir:'/root/autodl-tmp/project/nlp_transformers_tasks/text_classification/runs/test/'
# ckpt_filename_keyword:'checkpoint'
# ckpt_num:3
1200步之后
1600步之后