File size: 4,895 Bytes
40b9305 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | import os
import shutil
import glob
def get_dir_size(path):
total = 0
try:
for entry in os.scandir(path):
if entry.is_file():
total += entry.stat().st_size
elif entry.is_dir():
total += get_dir_size(entry.path)
except Exception:
pass
return total
def main():
project_dir = os.path.dirname(os.path.abspath(__file__))
experiments_dir = os.path.join(project_dir, "models", "CodeFormer", "experiments")
if not os.path.exists(experiments_dir):
print("Experiments directory not found.")
return
print("=== Disk Cleanup Script ===")
# Active run folder
active_folder = "20260708_201102_CodeFormer_stage3_custom"
active_path = os.path.join(experiments_dir, active_folder)
total_freed = 0
# 1. Delete old/abandoned experiment folders
for entry in os.scandir(experiments_dir):
if entry.is_dir() and entry.name != active_folder and entry.name != "pretrained_models":
dir_size = get_dir_size(entry.path)
print(f"Deleting old run: {entry.name} ({dir_size / 1024 / 1024 / 1024:.2f} GB)...")
try:
shutil.rmtree(entry.path)
total_freed += dir_size
print("Deleted.")
except Exception as e:
print(f"Error deleting {entry.name}: {e}")
# 2. Clean intermediate checkpoints in the active folder
if os.path.exists(active_path):
print(f"\nCleaning intermediate checkpoints in active run: {active_folder}...")
# Scan models directory
models_dir = os.path.join(active_path, "models")
if os.path.exists(models_dir):
pth_files = glob.glob(os.path.join(models_dir, "*.pth"))
# Find the highest iteration index in models
indices = []
for filepath in pth_files:
basename = os.path.basename(filepath)
if "latest" in basename:
continue
# format is net_g_XX.pth or net_d_XX.pth
try:
parts = basename.replace(".pth", "").split("_")
idx = int(parts[-1])
indices.append(idx)
except ValueError:
continue
if indices:
highest_idx = max(indices)
print(f"Active run highest iteration index: {highest_idx}")
# We want to keep:
# - files with "latest" in their name
# - files matching highest_idx
# - files matching highest_idx - 5 (as backup)
keep_indices = {highest_idx, highest_idx - 5}
for filepath in pth_files:
basename = os.path.basename(filepath)
if "latest" in basename:
continue
try:
parts = basename.replace(".pth", "").split("_")
idx = int(parts[-1])
if idx not in keep_indices:
file_size = os.path.getsize(filepath)
os.remove(filepath)
total_freed += file_size
except Exception as e:
print(f"Error removing {basename}: {e}")
# Scan training_states directory
states_dir = os.path.join(active_path, "training_states")
if os.path.exists(states_dir):
state_files = glob.glob(os.path.join(states_dir, "*.state"))
state_indices = []
for filepath in state_files:
basename = os.path.basename(filepath)
try:
idx = int(basename.replace(".state", ""))
state_indices.append(idx)
except ValueError:
continue
if state_indices:
highest_state_idx = max(state_indices)
keep_state_indices = {highest_state_idx, highest_state_idx - 5}
for filepath in state_files:
basename = os.path.basename(filepath)
try:
idx = int(basename.replace(".state", ""))
if idx not in keep_state_indices:
file_size = os.path.getsize(filepath)
os.remove(filepath)
total_freed += file_size
except Exception as e:
print(f"Error removing {basename}: {e}")
print(f"\nSUCCESS: Cleanup completed! Freed up {total_freed / 1024 / 1024 / 1024:.2f} GB of disk space.")
if __name__ == "__main__":
main()
|