Delete caffe/tensorflow checkpoint files with the non-last iteration. usage: python clean_up_ckpt.py --root_dir ~/workspace --ckpt_type tensorflow

clean_up_ckpt.py

import os
import argparse
import re
import numpy as np

parser = argparse.ArgumentParser()
parser.add_argument("--root_dir", help="The uppermost directory to clean up")
parser.add_argument("--ckpt_type", help="Ckpt type(tensorflow, caffe)")
args = parser.parse_args()

root_dir = '~/'
if args.root_dir:
    root_dir = args.root_dir
root_dir = os.path.expanduser(root_dir)

ckpt_type = 'tensorflow'
if args.ckpt_type:
    ckpt_type = args.ckpt_type

caffe_ckpt_pattern = re.compile('[0-9]+\.(caffemodel|solverstate)$')
tf_ckpt_pattern = re.compile('.ckpt-[0-9]+(.meta|.index|.data-[0-9]+-of-[0-9]+)?$')
int_pattern = re.compile('[0-9]+')

if ckpt_type == 'tensorflow':
    ckpt_pattern = tf_ckpt_pattern
elif ckpt_type == 'caffe':
    ckpt_pattern = caffe_ckpt_pattern
else:
    raise Exception('Wrong checkpoint type: %s' % ckpt_type)


def get_response(question, default='y'):
    resp = raw_input('%s[%s]: ' % (question, default))
    if resp == '':
        resp = default
    return resp

def get_ckpt_w_iter(files, pattern):
    if not files or len(files) == 0:
        return None
    ret = []
    for fname in files:
        match = pattern.search(fname)
        if match:
            match_str = match.group()
            i = int(int_pattern.search(match_str).group())
            ret.append((fname, i))
    return ret


# Iterating directories, find whether dir has ckpt files
for root, dirs, files in os.walk(root_dir):
    ckpts = get_ckpt_w_iter(files, ckpt_pattern)
    if not ckpts or len(ckpts) == 0:
        continue

    # Find the maximum iter# and whether there are ckpts to delete
    max_iter = np.max([temp[1] for temp in ckpts])
    max_ckpts = [temp for temp in ckpts if temp[1] == max_iter]
    other_ckpts = [temp for temp in ckpts if temp[1] < max_iter]
    if len(other_ckpts) == 0:
        continue

    print root
    print('\tckpt with max iter: %s' % str(max_ckpts))
    print('\tother ckpt: %s' % str(other_ckpts))

    # Delete non-last ckpt files
    resp = get_response('Delete %d other ckpts?[y/n]' % len(other_ckpts))
    if resp == 'y':
        for fname, i in other_ckpts:
            print('\tdeleting %s....' % fname)
            os.remove(os.path.join(root, fname))
        print ' '