Skip to content

Commit cf3fb24

Browse files
committed
add clean checkpoint
1 parent 192f9a5 commit cf3fb24

File tree

1 file changed

+13
-0
lines changed

1 file changed

+13
-0
lines changed

python/paddle/fluid/io.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -529,6 +529,19 @@ def load_checkpoint(executor, checkpoint_dir=None, main_program=None):
529529
filename=None)
530530

531531

532+
def clean_checkpoint(checkpoint_dir, delete_dir=False):
533+
"""
534+
clean the checkpoint dir, when the train exits normally, the trainer will call clean_checkpoint to delete checkpoint directory saved before.
535+
delete_dir only works when the directory is empty, otherwise, OSError is raised.
536+
"""
537+
if checkpoint_dir is None:
538+
checkpoint_dir = os.getcwd()
539+
_lru_delete(checkpoint_dir, max_num_checkpoints=0)
540+
541+
if delete_dir and not os.listdir(checkpoint_dir):
542+
os.rmdir(checkpoint_dir)
543+
544+
532545
def _get_serial_dir(serial, checkpoint_dir):
533546
serial_folder = CHECKPOINT_PREFIX + CHECKPOINT_SEPARATOR + str(serial)
534547
return os.path.join(checkpoint_dir, serial_folder)

0 commit comments

Comments
 (0)