File size: 2,869 Bytes
31fc7e1 |
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 |
import os
import json
from datetime import datetime
def create_run_directory(base_dir='runs', prefix='run', suffix='', parent_dir=None):
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
dir_name = f"{prefix}_{timestamp}{suffix}"
if parent_dir:
run_dir = os.path.join(parent_dir, dir_name)
else:
run_dir = os.path.join(base_dir, dir_name)
os.makedirs(run_dir, exist_ok=True)
return run_dir
# Find the most recent run directory
def get_latest_run_dir(base_dir='runs', include_hyperparam=True):
all_dirs = []
for d in os.listdir(base_dir):
if d.startswith('run_'):
full_path = os.path.join(base_dir, d)
all_dirs.append(full_path)
if d.endswith('_hyperparam') and include_hyperparam:
# If it's a hyperparam directory, add its trial subdirectories
trial_dirs = [os.path.join(full_path, td) for td in os.listdir(full_path) if td.startswith('trial_')]
all_dirs.extend(trial_dirs)
if not all_dirs:
raise ValueError(f"No run directories found in {base_dir}")
# Sort directories by timestamp in the directory name
return max(all_dirs, key=get_dir_timestamp)
def get_run_file(filename, run_dir=None, required=True):
"""Get a file from a run directory
Args:
filename: Name of file to get (e.g., 'best_model.pth', 'config.json')
run_dir: Run directory path (uses latest if None)
required: Whether to raise an error if file not found
Returns:
str: Path to the file
dict: Loaded JSON data if file ends with .json
"""
if run_dir is None:
run_dir = get_latest_run_dir()
file_path = os.path.join(run_dir, filename)
if not os.path.exists(file_path):
if required:
raise FileNotFoundError(f"{filename} not found in {run_dir}")
return None
# Load JSON files automatically
if filename.endswith('.json'):
with open(file_path, 'r') as f:
return json.load(f)
return file_path
def get_latest_model_path(run_dir=None):
"""Get path to best_model.pth"""
return get_run_file('best_model.pth', run_dir)
def get_config(run_dir=None):
"""Get config from run directory"""
return get_run_file('config.json', run_dir)
# Helper function to parse directory name and get timestamp
def get_dir_timestamp(dir_path):
dir_name = os.path.basename(dir_path)
try:
# Extract timestamp from directory name
timestamp_str = dir_name.split('_')[1] # Assumes format is always prefix_timestamp or prefix_timestamp_suffix
return datetime.strptime(timestamp_str, "%Y%m%d-%H%M%S")
except (IndexError, ValueError):
# If parsing fails, return the earliest possible date
return datetime.min
|