clean_nets.py

"""
Clean trained nets (*/model-epoch_#.pth) whose epoch is neither the largest nor a multiple of 50
"""
import sys
import os

sys.path.append(os.path.abspath(sys.path[0] + '/../'))


if __name__ == "__main__":
    for dirpath, _, filenames in os.walk('../data'):
        epoch_list = [int(filename[12:-4]) for filename in filenames
                        if filename.startswith("model-epoch_")]
        if len(epoch_list) <= 1:
            continue
        epoch_list.sort()
        for epoch in epoch_list[:-1]:
            if epoch % 50 != 0:
                file_to_del = f"{dirpath}/model-epoch_{epoch}.pth"
                print(f"Clean model file: {file_to_del}")
                os.remove(file_to_del)