模型训练时维护checkpoints文件夹的数量

起因

在训练模型时，要每隔一定步数 要验证一次，如果验证指标更好了，则要保存对应的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步之后