#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Created on Tue Mar 25 10:37:33 2025 @author: louis """ import os import glob import shutil import torch import yaml def get_all_version_folders_containing_checkpoint( base_path="./lightning_logs/taslp", ): checkpoint_paths = glob.glob( os.path.join(base_path, "**", "checkpoints"), recursive=True ) return [ os.path.abspath(os.path.join(cp, os.pardir)) for cp in checkpoint_paths ] def get_latest_valid_version(base_path): # valid if contains ckpt folder version_paths = os.listdir(base_path) valid_version_paths = [] for version_path in version_paths: abs_version_path = os.path.abspath( os.path.join(base_path, version_path) ) if "checkpoints" in os.listdir(abs_version_path): valid_version_paths.append(abs_version_path) if not valid_version_paths: return "" return max(valid_version_paths, key=os.path.getmtime) def get_latest_checkpoint(base_dir): if base_dir.endswith(".ckpt"): base_dir = os.path.dirname(base_dir) if not "version_" in base_dir: base_dir = get_latest_valid_version(base_dir) if not base_dir.endswith("checkpoints"): base_dir = os.path.join(base_dir, "checkpoints") all_checkpoints = [ os.path.abspath(os.path.join(base_dir, ckpt_path)) for ckpt_path in os.listdir(base_dir) ] return max(all_checkpoints, key=os.path.getmtime) def get_best_checkpoint(base_dir, monitor_mode): latest_checkpoint = get_latest_checkpoint(base_dir) d = torch.load(latest_checkpoint, weights_only=False, map_location="cpu") model_checkpoint_callbacks = { k: v for k, v in d["callbacks"].items() if "ModelCheckpoint" in k } if len(model_checkpoint_callbacks) > 1: raise RuntimeError("Not that many callbacks expected") model_checkpoint_callback = next(iter(model_checkpoint_callbacks.values())) best_k_models = model_checkpoint_callback["best_k_models"] if len(best_k_models) == 0: print( "No checkpoint dict ranked by monitor found, defaulting to latest checkpoint" ) best_ckpt = latest_checkpoint else: best_ckpt = monitor_mode(best_k_models, key=best_k_models.get) print(f"best_monitor_value: {best_k_models[best_ckpt].item()}") # correct if latest_checkpoint is computed from a sshfs-mounted dir best_ckpt_in_latest_ckpt_dir = os.path.join( os.path.dirname(latest_checkpoint), os.path.basename(best_ckpt) ) return best_ckpt_in_latest_ckpt_dir def move_all_except_lastest_version_containing_checkpoint_to_trash( base_path="./lightning_logs/taslp", trash_path="./lightning_logs/trash", dry_run=True, clear_all_except_config_and_best_checkpoint: bool = False, clear_all_except_config_and_best_checkpoint_monitor_mode=max, delete_test_versions: bool = True, backup_all_ckpts: bool = True, # false for less disk space usage but very dagerous ): # find all "version_*" dirs (in all loggers) version_paths = sorted( glob.glob(os.path.join(base_path, "**", "version_*"), recursive=True) ) # We handle test versions differently if delete_test_versions: for vp in version_paths: if "test" in vp: logger_path = os.path.abspath(os.path.join(vp, os.pardir)) logger_name = os.path.relpath(logger_path, base_path) source_dir = vp if os.path.isdir(source_dir): trash_dest = ( os.path.join(trash_path, logger_name) + "_test" ) print(f"moving {source_dir} to {trash_dest}") if not dry_run: try: os.makedirs(trash_dest, exist_ok=True) shutil.move(source_dir, trash_dest) except Exception as e: print(e) breakpoint() version_paths = sorted( glob.glob(os.path.join(base_path, "**", "version_*"), recursive=True) ) # Extract all logger_paths logger_paths = set( os.path.abspath(os.path.join(vp, os.pardir)) for vp in version_paths ) for logger_path in logger_paths: # get logger_name logger_name = os.path.relpath(logger_path, base_path) trash_dest = os.path.join(trash_path, logger_name) # get latest version latest_valid_version = get_latest_valid_version(base_path=logger_path) for version in os.listdir(logger_path): source_dir = os.path.join(logger_path, version) if version != os.path.basename(latest_valid_version): if not dry_run: os.makedirs(trash_dest, exist_ok=True) print(f"moving {source_dir} to {trash_dest}") if not dry_run: try: shutil.move(source_dir, trash_dest) except Exception as e: print(e) breakpoint() else: if clear_all_except_config_and_best_checkpoint: if backup_all_ckpts: print(f"copying {source_dir} to {trash_dest}") if not dry_run: try: shutil.copytree( source_dir, trash_dest, copy_function=shutil.copy2, ) except Exception as e: print(e) breakpoint() all_files = glob.glob( os.path.join(source_dir, "**"), recursive=True ) best_checkpoint = get_best_checkpoint( os.path.join(source_dir, "checkpoints"), monitor_mode=clear_all_except_config_and_best_checkpoint_monitor_mode, ) for file in all_files: if ( not file.endswith(".yaml") and os.path.basename(file) != os.path.basename(best_checkpoint) and not os.path.isdir(file) ): print(f"Deleting {file}") if not dry_run: os.remove(file) def instantiate_model_only( config_path, ckpt_path: str | None = None, remove_reverb_model_and_joint_loss=True, ): # we do local import to avoid circular deps from cli import MyCli from datasets import AudioDatasetConvolvedWithRirDatasetDataModule from model.joint_model import JointModel # https://github.com/Lightning-AI/pytorch-lightning/issues/17447 # latest_valid_config= config = yaml.load(open(config_path, "r"), Loader=yaml.FullLoader) args_dict = { key: value for key, value in config.items() if key in ("model", "seed_everything") } | {"data": {"class_path": "datasets.NoDataModule"}} if remove_reverb_model_and_joint_loss: args_dict["model"]["reverb_model"] = None args_dict["model"]["joint_loss_module"] = None if ckpt_path is None: print("Using best checkpoint") ckpt_path = get_best_checkpoint( os.path.dirname(config_path), monitor_mode=max ) args_dict["model"]["speech_model_ckpt_path"] = ckpt_path cli = MyCli( model_class=JointModel, datamodule_class=AudioDatasetConvolvedWithRirDatasetDataModule, subclass_mode_model=False, subclass_mode_data=True, run=False, args=args_dict, ) model = cli.model return model