From 150c2a725301b7de374317d1b6d8e632e33f169e Mon Sep 17 00:00:00 2001 From: YannAhlgrim Date: Tue, 2 Jun 2026 12:00:12 +0200 Subject: [PATCH] fix: do not run 4 times and log new params --- configs/grids/lp_grid.yaml | 4 +- src/train_supervised.py | 699 +++++++++++++++++++------------------ src/utils/logging.py | 4 + 3 files changed, 360 insertions(+), 347 deletions(-) diff --git a/configs/grids/lp_grid.yaml b/configs/grids/lp_grid.yaml index 51027ae..af93730 100644 --- a/configs/grids/lp_grid.yaml +++ b/configs/grids/lp_grid.yaml @@ -10,8 +10,8 @@ grid: optimization.momentum: [0.9, 0.99, 0.5] optimization.lr_schedule: [cosine, step] data.batch_size: [32, 64, 128, 256] - meta.representation_type: [last_avgpool] - meta.head_type: [linear] + meta.representation_type: [last_avgpool, last4_avgpool_concat] + meta.head_type: [linear, bn_linear] launch: folder: submitit_logs/ diff --git a/src/train_supervised.py b/src/train_supervised.py index 0e6de44..eb3eb20 100644 --- a/src/train_supervised.py +++ b/src/train_supervised.py @@ -59,6 +59,42 @@ def distributed_sum(value, device): return float(tensor.item()) +def _find_metric(metrics, key): + if isinstance(metrics, dict): + if key in metrics: + return metrics[key] + for value in metrics.values(): + found = _find_metric(value, key) + if found is not None: + return found + if isinstance(metrics, list): + for item in metrics: + found = _find_metric(item, key) + if found is not None: + return found + return None + + +def _collect_eval_rows(eval_root, metric_key): + rows = [] + for dirpath, _, filenames in os.walk(eval_root): + for fname in filenames: + if not fname.endswith("_metrics.json"): + continue + path = os.path.join(dirpath, fname) + try: + with open(path, "r") as f: + metrics = json.load(f) + except (OSError, json.JSONDecodeError): + continue + value = _find_metric(metrics, metric_key) + if value is None: + continue + run_name = os.path.basename(os.path.dirname(path)) + rows.append((float(value), run_name, path)) + return rows + + class EarlyStopping: def __init__( self, @@ -178,14 +214,93 @@ def main(args, resume_preempt=False): v_args = args["validation"] es_args = o_args["early_stopping"] - root_folder = resolve_log_dir(args, stage="train") + folder = resolve_log_dir(args, stage="train") tag = l_args["write_tag"] - with open(os.path.join(root_folder, "params-supervised.yaml"), "w") as f: + with open(os.path.join(folder, "params-supervised.yaml"), "w") as f: yaml.dump(args, f) - with open(os.path.join(root_folder, "params.yaml"), "w") as f: + with open(os.path.join(folder, "params.yaml"), "w") as f: yaml.dump(args, f) + save_path = os.path.join(folder, f"{tag}" + "-ep{epoch}.pth.tar") + latest_path = os.path.join(folder, f"{tag}-latest.pth.tar") + best_path = os.path.join(folder, f"{tag}-best.pth.tar") + log_file = os.path.join(folder, f"{tag}_r{rank}.csv") + + csv_logger = CSVLogger( + log_file, + ("%d", "epoch"), + ("%.6f", "train_loss"), + ("%.6f", "val_loss"), + ("%.6f", "val_acc"), + ("%.6e", "lr"), + ("%.6f", "best_val_loss"), + ("%d", "best_epoch"), + ("%d", "early_stop"), + ) + + encoder, _ = init_model( + device=device, + patch_size=mk_args["patch_size"], + crop_size=d_args["crop_size"], + model_name=m_args["model_name"], + ) + + representation_type = m_args.get("representation_type", "last_avgpool") + head_type = m_args.get("head_type", "linear") + model = ViTClassifier( + encoder, + m_args["num_classes"], + m_args["embed_dim"], + representation_type=representation_type, + head_type=head_type, + ).to(device) + + if o_args["freeze_weights"]: + logger.info("Freezing encoder weights (Linear Probing mode)") + for param in model.encoder.parameters(): + param.requires_grad = False + model.encoder.eval() + else: + logger.info("Training full model (Fine-tuning mode)") + + params = [p for p in model.parameters() if p.requires_grad] + optimizer_name = o_args["optimizer"].lower() + if optimizer_name == "adamw": + optimizer = torch.optim.AdamW( + params, lr=o_args["lr"], weight_decay=o_args["weight_decay"] + ) + elif optimizer_name == "lars": + optimizer = LARS( + params, + lr=o_args["lr"], + weight_decay=o_args["weight_decay"], + momentum=o_args.get("momentum", 0.9), + eta=o_args.get("lars_eta", 0.001), + eps=o_args.get("lars_eps", 1e-8), + exclude_bias_and_norm=o_args.get("lars_exclude_bias_and_norm", True), + ) + else: + optimizer = torch.optim.SGD( + params, + lr=o_args["lr"], + momentum=o_args.get("momentum", 0.9), + weight_decay=o_args["weight_decay"], + ) + + scheduler = None + lr_schedule = o_args.get("lr_schedule", "cosine").lower() + if lr_schedule == "step": + scheduler = torch.optim.lr_scheduler.MultiStepLR( + optimizer, + milestones=o_args.get("step_milestones", [15, 30, 45]), + gamma=o_args.get("step_gamma", 0.1), + ) + elif lr_schedule == "cosine": + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( + optimizer, T_max=o_args["epochs"], eta_min=o_args["final_lr"] + ) + train_transform = make_transforms( crop_size=d_args["crop_size"], crop_scale=tuple(d_args["crop_scale"]), @@ -225,369 +340,263 @@ def main(args, resume_preempt=False): drop_last=False, ) - combinations = [ - ("last_avgpool", "linear"), - ("last_avgpool", "bn_linear"), - ("last4_avgpool_concat", "linear"), - ("last4_avgpool_concat", "bn_linear"), - ] - criterion = nn.CrossEntropyLoss().to(device) + model = DistributedDataParallel(model, device_ids=[torch.cuda.current_device()]) + eval_every = int(v_args["eval_every"]) - combo_results = [] - for representation_type, head_type in combinations: - combo_name = f"{representation_type}__{head_type}" - combo_folder = os.path.join(root_folder, combo_name) - os.makedirs(combo_folder, exist_ok=True) + early_stopper = EarlyStopping( + enabled=es_args["enabled"], + patience=es_args["patience"], + min_delta=es_args["min_delta"], + min_epochs=es_args["min_epochs"], + restore_best_weights=es_args["restore_best_weights"], + ) - combo_args = yaml.safe_load(yaml.dump(args)) - combo_args.setdefault("meta", {})["representation_type"] = representation_type - combo_args.setdefault("meta", {})["head_type"] = head_type + start_epoch = 0 - with open(os.path.join(combo_folder, "params-supervised.yaml"), "w") as f: - yaml.dump(combo_args, f) - with open(os.path.join(combo_folder, "params.yaml"), "w") as f: - yaml.dump(combo_args, f) + checkpoint_to_load = None + resuming_interrupted = False + if os.path.exists(latest_path): + checkpoint_to_load = latest_path + resuming_interrupted = True + elif m_args["load_checkpoint"]: + r_file = m_args["read_checkpoint"] + checkpoint_folder = m_args["checkpoint_folder"] + if os.path.isabs(r_file): + checkpoint_to_load = r_file + else: + checkpoint_to_load = os.path.join(checkpoint_folder, r_file) - save_path = os.path.join(combo_folder, f"{tag}" + "-ep{epoch}.pth.tar") - latest_path = os.path.join(combo_folder, f"{tag}-latest.pth.tar") - best_path = os.path.join(combo_folder, f"{tag}-best.pth.tar") - log_file = os.path.join(combo_folder, f"{tag}_r{rank}.csv") + if checkpoint_to_load is not None and not os.path.exists(checkpoint_to_load): + raise FileNotFoundError(f"Checkpoint not found: {checkpoint_to_load}") - csv_logger = CSVLogger( - log_file, - ("%d", "epoch"), - ("%.6f", "train_loss"), - ("%.6f", "val_loss"), - ("%.6f", "val_acc"), - ("%.6e", "lr"), - ("%.6f", "best_val_loss"), - ("%d", "best_epoch"), - ("%d", "early_stop"), - ) - - encoder, _ = init_model( - device=device, - patch_size=mk_args["patch_size"], - crop_size=d_args["crop_size"], - model_name=m_args["model_name"], - ) - - model = ViTClassifier( - encoder, - m_args["num_classes"], - m_args["embed_dim"], - representation_type=representation_type, - head_type=head_type, - ).to(device) - - if o_args["freeze_weights"]: + if checkpoint_to_load and os.path.exists(checkpoint_to_load): + checkpoint = torch.load(checkpoint_to_load, map_location="cpu") + if resuming_interrupted and "model" in checkpoint: + model.module.load_state_dict(checkpoint["model"]) + if "opt" in checkpoint: + optimizer.load_state_dict(checkpoint["opt"]) + if scheduler is not None and "scheduler" in checkpoint: + scheduler.load_state_dict(checkpoint["scheduler"]) + start_epoch = int(checkpoint.get("epoch", 0)) + early_stopper.load_state_dict(checkpoint.get("early_stopping", {})) logger.info( - f"[{combo_name}] Freezing encoder weights (Linear Probing mode)" - ) - for param in model.encoder.parameters(): - param.requires_grad = False - model.encoder.eval() - else: - logger.info(f"[{combo_name}] Training full model (Fine-tuning mode)") - - params = [p for p in model.parameters() if p.requires_grad] - optimizer_name = o_args["optimizer"].lower() - if optimizer_name == "adamw": - optimizer = torch.optim.AdamW( - params, lr=o_args["lr"], weight_decay=o_args["weight_decay"] - ) - elif optimizer_name == "lars": - optimizer = LARS( - params, - lr=o_args["lr"], - weight_decay=o_args["weight_decay"], - momentum=o_args.get("momentum", 0.9), - eta=o_args.get("lars_eta", 0.001), - eps=o_args.get("lars_eps", 1e-8), - exclude_bias_and_norm=o_args.get("lars_exclude_bias_and_norm", True), + f"Resuming training from {checkpoint_to_load} at epoch {start_epoch}" ) else: - optimizer = torch.optim.SGD( - params, - lr=o_args["lr"], - momentum=o_args.get("momentum", 0.9), - weight_decay=o_args["weight_decay"], + encoder_state = checkpoint.get("encoder") + if encoder_state is None and "model" in checkpoint: + encoder_state = { + k.replace("encoder.", "", 1): v + for k, v in checkpoint["model"].items() + if k.startswith("encoder.") + } + if encoder_state is None: + raise KeyError( + f"No encoder weights found in checkpoint: {checkpoint_to_load}" + ) + encoder_state = strip_module_prefix(encoder_state) + msg = model.module.encoder.load_state_dict(encoder_state, strict=False) + logger.info( + f"Loaded pre-trained encoder from {checkpoint_to_load} with msg: {msg}" ) - scheduler = None - lr_schedule = o_args.get("lr_schedule", "cosine").lower() - if lr_schedule == "step": - scheduler = torch.optim.lr_scheduler.MultiStepLR( - optimizer, - milestones=o_args.get("step_milestones", [15, 30, 45]), - gamma=o_args.get("step_gamma", 0.1), - ) - elif lr_schedule == "cosine": - scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( - optimizer, T_max=o_args["epochs"], eta_min=o_args["final_lr"] - ) - - model = DistributedDataParallel(model, device_ids=[torch.cuda.current_device()]) - - early_stopper = EarlyStopping( - enabled=es_args["enabled"], - patience=es_args["patience"], - min_delta=es_args["min_delta"], - min_epochs=es_args["min_epochs"], - restore_best_weights=es_args["restore_best_weights"], - ) - - start_epoch = 0 - - checkpoint_to_load = None - resuming_interrupted = False - if os.path.exists(latest_path): - checkpoint_to_load = latest_path - resuming_interrupted = True - elif m_args["load_checkpoint"]: - r_file = m_args["read_checkpoint"] - checkpoint_folder = m_args["checkpoint_folder"] - if os.path.isabs(r_file): - checkpoint_to_load = r_file - else: - checkpoint_to_load = os.path.join(checkpoint_folder, r_file) - - if checkpoint_to_load is not None and not os.path.exists(checkpoint_to_load): - raise FileNotFoundError(f"Checkpoint not found: {checkpoint_to_load}") - - if checkpoint_to_load and os.path.exists(checkpoint_to_load): - checkpoint = torch.load(checkpoint_to_load, map_location="cpu") - if resuming_interrupted and "model" in checkpoint: - model.module.load_state_dict(checkpoint["model"]) - if "opt" in checkpoint: - optimizer.load_state_dict(checkpoint["opt"]) - if scheduler is not None and "scheduler" in checkpoint: - scheduler.load_state_dict(checkpoint["scheduler"]) - start_epoch = int(checkpoint.get("epoch", 0)) - early_stopper.load_state_dict(checkpoint.get("early_stopping", {})) - logger.info( - f"[{combo_name}] Resuming training from {checkpoint_to_load} at epoch {start_epoch}" - ) - else: - encoder_state = checkpoint.get("encoder") - if encoder_state is None and "model" in checkpoint: - encoder_state = { - k.replace("encoder.", "", 1): v - for k, v in checkpoint["model"].items() - if k.startswith("encoder.") - } - if encoder_state is None: - raise KeyError( - f"No encoder weights found in checkpoint: {checkpoint_to_load}" - ) - encoder_state = strip_module_prefix(encoder_state) - msg = model.module.encoder.load_state_dict(encoder_state, strict=False) - logger.info( - f"[{combo_name}] Loaded pre-trained encoder from {checkpoint_to_load} with msg: {msg}" - ) - - def save_checkpoint(epoch, train_loss, val_loss, val_acc, is_best=False): - save_dict = { - "model": model.module.state_dict(), - "opt": optimizer.state_dict(), - "scheduler": None if scheduler is None else scheduler.state_dict(), - "epoch": epoch, - "train_loss": train_loss, - "val_loss": val_loss, - "val_acc": val_acc, - "args": combo_args, - "early_stopping": early_stopper.state_dict(), - } - if rank == 0: - torch.save(save_dict, latest_path) - if epoch % checkpoint_freq == 0: - torch.save(save_dict, save_path.format(epoch=epoch)) - if is_best: - torch.save(save_dict, best_path) - - for epoch in range(start_epoch, o_args["epochs"]): - train_sampler.set_epoch(epoch) - val_sampler.set_epoch(epoch) - - model.train() - if o_args["freeze_weights"]: - model.module.encoder.eval() - - loss_meter = AverageMeter() - - for itr, (imgs, labels) in enumerate(train_loader): - imgs = imgs.to(device, non_blocking=True) - labels = labels.to(device, non_blocking=True) - - with torch.cuda.amp.autocast( - enabled=m_args["use_bfloat16"], dtype=torch.bfloat16 - ): - outputs = model(imgs) - loss = criterion(outputs, labels) - - optimizer.zero_grad() - loss.backward() - optimizer.step() - - loss_meter.update(loss.item(), n=labels.size(0)) - - if itr % log_freq == 0 and rank == 0: - logger.info( - f"[{combo_name}] Epoch {epoch + 1} [{itr}/{len(train_loader)}] Train Loss: {loss_meter.avg:.4f}" - ) - - train_loss = distributed_average(loss_meter.avg, device) - - do_eval = ((epoch + 1) % eval_every == 0) or ( - epoch + 1 == o_args["epochs"] - ) - if do_eval: - val_loss, val_acc = evaluate( - model=model, - loader=val_loader, - criterion=criterion, - device=device, - use_bfloat16=m_args["use_bfloat16"], - ) - is_best, should_stop = early_stopper.step( - epoch + 1, val_loss, model.module - ) - else: - val_loss = float("nan") - val_acc = float("nan") - is_best, should_stop = False, False - - if scheduler is not None: - scheduler.step() - - if rank == 0: - logger.info( - f"[{combo_name}] Epoch {epoch + 1} done | train_loss={train_loss:.6f} val_loss={val_loss:.6f} val_acc={val_acc:.6f} best_val_loss={early_stopper.best_metric:.6f}" - ) - - csv_logger.log( - epoch + 1, - train_loss, - val_loss, - val_acc, - optimizer.param_groups[0]["lr"], - early_stopper.best_metric, - early_stopper.best_epoch, - int(should_stop), - ) - - save_checkpoint(epoch + 1, train_loss, val_loss, val_acc, is_best=is_best) - - stop_tensor = torch.tensor([int(should_stop)], device=device) - if ( - dist.is_available() - and dist.is_initialized() - and dist.get_world_size() > 1 - ): - dist.broadcast(stop_tensor, src=0) - - if bool(stop_tensor.item()): - if rank == 0: - logger.info( - f"[{combo_name}] Early stopping at epoch {epoch + 1}. Best val_loss={early_stopper.best_metric:.6f} @ epoch {early_stopper.best_epoch}" - ) - break - - if early_stopper.enabled and early_stopper.restore_best_weights: - if rank == 0: - logger.info(f"[{combo_name}] Restoring best model weights before exit") - early_stopper.restore(model.module, device) - if rank == 0: - torch.save( - { - "model": model.module.state_dict(), - "epoch": early_stopper.best_epoch, - "val_loss": early_stopper.best_metric, - "args": combo_args, - }, - best_path, - ) - - combo_result = { - "representation_type": representation_type, - "head_type": head_type, - "combo_name": combo_name, - "best_val_loss": float(early_stopper.best_metric), - "best_epoch": int(early_stopper.best_epoch), - "best_checkpoint": best_path, - "train_folder": combo_folder, + def save_checkpoint(epoch, train_loss, val_loss, val_acc, is_best=False): + save_dict = { + "model": model.module.state_dict(), + "opt": optimizer.state_dict(), + "scheduler": None if scheduler is None else scheduler.state_dict(), + "epoch": epoch, + "train_loss": train_loss, + "val_loss": val_loss, + "val_acc": val_acc, + "args": args, + "early_stopping": early_stopper.state_dict(), } + if rank == 0: + torch.save(save_dict, latest_path) + if epoch % checkpoint_freq == 0: + torch.save(save_dict, save_path.format(epoch=epoch)) + if is_best: + torch.save(save_dict, best_path) + + for epoch in range(start_epoch, o_args["epochs"]): + train_sampler.set_epoch(epoch) + val_sampler.set_epoch(epoch) + + model.train() + if o_args["freeze_weights"]: + model.module.encoder.eval() + + loss_meter = AverageMeter() + + for itr, (imgs, labels) in enumerate(train_loader): + imgs = imgs.to(device, non_blocking=True) + labels = labels.to(device, non_blocking=True) + + with torch.cuda.amp.autocast( + enabled=m_args["use_bfloat16"], dtype=torch.bfloat16 + ): + outputs = model(imgs) + loss = criterion(outputs, labels) + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + loss_meter.update(loss.item(), n=labels.size(0)) + + if itr % log_freq == 0 and rank == 0: + logger.info( + f"Epoch {epoch + 1} [{itr}/{len(train_loader)}] Train Loss: {loss_meter.avg:.4f}" + ) + + train_loss = distributed_average(loss_meter.avg, device) + + do_eval = ((epoch + 1) % eval_every == 0) or (epoch + 1 == o_args["epochs"]) + if do_eval: + val_loss, val_acc = evaluate( + model=model, + loader=val_loader, + criterion=criterion, + device=device, + use_bfloat16=m_args["use_bfloat16"], + ) + is_best, should_stop = early_stopper.step(epoch + 1, val_loss, model.module) + else: + val_loss = float("nan") + val_acc = float("nan") + is_best, should_stop = False, False + + if scheduler is not None: + scheduler.step() if rank == 0: - eval_args = { - "meta": { - "seed": m_args.get("seed", _GLOBAL_SEED), - "model_name": m_args["model_name"], - "embed_dim": m_args["embed_dim"], - "num_classes": m_args["num_classes"], - "patch_size": mk_args.get( - "patch_size", m_args.get("patch_size", 16) - ), - "crop_size": d_args.get("crop_size", m_args.get("crop_size", 224)), - "use_bfloat16": m_args.get("use_bfloat16", True), - "representation_type": representation_type, - "head_type": head_type, - "checkpoint_path": best_path, - "force_single_process": True, - }, - "data": { - "batch_size": d_args.get("batch_size", 128), - "root_path": d_args.get("root_path", "./wilds_data"), - "num_workers": d_args.get("num_workers", 8), - "pin_mem": d_args.get("pin_mem", True), - "split": "test", - "download": True, - }, - "logging": { - "write_tag": "iwildcam_test", - "auto_folder": True, - }, - } - eval_result = eval_wilds_main(args=eval_args) - eval_folder = eval_args.get("logging", {}).get("folder") - supervised_params = os.path.join(combo_folder, "params-supervised.yaml") - if eval_folder and os.path.exists(supervised_params): - try: - shutil.copy2( - supervised_params, - os.path.join(eval_folder, "params-supervised.yaml"), - ) - shutil.copy2( - supervised_params, - os.path.join(eval_folder, "params.yaml"), - ) - except OSError: - logger.warning("Could not copy supervised params to eval folder") + logger.info( + f"Epoch {epoch + 1} done | train_loss={train_loss:.6f} val_loss={val_loss:.6f} val_acc={val_acc:.6f} best_val_loss={early_stopper.best_metric:.6f}" + ) - combo_result["eval"] = eval_result + csv_logger.log( + epoch + 1, + train_loss, + val_loss, + val_acc, + optimizer.param_groups[0]["lr"], + early_stopper.best_metric, + early_stopper.best_epoch, + int(should_stop), + ) - combo_results.append(combo_result) + save_checkpoint(epoch + 1, train_loss, val_loss, val_acc, is_best=is_best) - if ( - dist.is_available() - and dist.is_initialized() - and dist.get_world_size() > 1 - ): - dist.barrier() + stop_tensor = torch.tensor([int(should_stop)], device=device) + if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: + dist.broadcast(stop_tensor, src=0) + + if bool(stop_tensor.item()): + if rank == 0: + logger.info( + f"Early stopping at epoch {epoch + 1}. Best val_loss={early_stopper.best_metric:.6f} @ epoch {early_stopper.best_epoch}" + ) + break + + if early_stopper.enabled and early_stopper.restore_best_weights: + if rank == 0: + logger.info("Restoring best model weights before exit") + early_stopper.restore(model.module, device) + if rank == 0: + torch.save( + { + "model": model.module.state_dict(), + "epoch": early_stopper.best_epoch, + "val_loss": early_stopper.best_metric, + "args": args, + }, + best_path, + ) if rank == 0: - best_combo = min(combo_results, key=lambda r: r["best_val_loss"]) - summary = { - "protocol": "ijepa_linear_probe_best_of_4", - "combinations": combo_results, - "best_by_val_loss": best_combo, + eval_args = { + "meta": { + "seed": m_args.get("seed", _GLOBAL_SEED), + "model_name": m_args["model_name"], + "embed_dim": m_args["embed_dim"], + "num_classes": m_args["num_classes"], + "patch_size": mk_args.get("patch_size", m_args.get("patch_size", 16)), + "crop_size": d_args.get("crop_size", m_args.get("crop_size", 224)), + "use_bfloat16": m_args.get("use_bfloat16", True), + "representation_type": representation_type, + "head_type": head_type, + "checkpoint_path": best_path, + "force_single_process": True, + }, + "data": { + "batch_size": d_args.get("batch_size", 128), + "root_path": d_args.get("root_path", "./wilds_data"), + "num_workers": d_args.get("num_workers", 8), + "pin_mem": d_args.get("pin_mem", True), + "split": "test", + "download": True, + }, + "logging": { + "write_tag": "iwildcam_test", + "auto_folder": True, + }, } - summary_path = os.path.join(root_folder, "ijepa_linear_probe_summary.json") - with open(summary_path, "w") as f: - json.dump(summary, f, indent=2, sort_keys=True) - logger.info(f"Saved IJEPA LP summary to {summary_path}") + eval_result = eval_wilds_main(args=eval_args) + + eval_folder = eval_args.get("logging", {}).get("folder") + if eval_folder: + try: + params_out = yaml.safe_load(yaml.dump(args)) + params_out.setdefault("meta", {})["representation_type"] = representation_type + params_out.setdefault("meta", {})["head_type"] = head_type + metric_key = m_args.get("selection_metric", "macro_f1") + eval_root = os.path.join("experiment_logs", "eval-wilds") + ranking = { + "metric_key": metric_key, + "this_run_name": os.path.basename(os.path.normpath(eval_folder)), + "this_is_best": None, + "this_metric": None, + "best_run_name": None, + "best_metric": None, + } + if os.path.isdir(eval_root): + rows = _collect_eval_rows(eval_root, metric_key) + if rows: + rows.sort(key=lambda r: r[0], reverse=True) + best_value, best_run_name, _ = rows[0] + this_run_name = ranking["this_run_name"] + this_value = None + for value, run_name, _ in rows: + if run_name == this_run_name: + this_value = value + break + ranking["this_metric"] = this_value + ranking["best_run_name"] = best_run_name + ranking["best_metric"] = best_value + if this_value is not None: + ranking["this_is_best"] = this_run_name == best_run_name + params_out["results"] = { + "best_val_loss": float(early_stopper.best_metric), + "best_epoch": int(early_stopper.best_epoch), + "best_checkpoint": best_path, + "eval_metrics": eval_result.get("metrics") if eval_result else None, + "ranking": ranking, + } + with open(os.path.join(eval_folder, "params-supervised.yaml"), "w") as f: + yaml.dump(params_out, f) + with open(os.path.join(eval_folder, "params.yaml"), "w") as f: + yaml.dump(params_out, f) + except OSError: + logger.warning("Could not write supervised params to eval folder") + + if os.path.exists(folder): + try: + shutil.rmtree(folder) + except OSError: + logger.warning("Could not remove supervised run folder") if __name__ == "__main__": diff --git a/src/utils/logging.py b/src/utils/logging.py index e54d363..5caab7c 100644 --- a/src/utils/logging.py +++ b/src/utils/logging.py @@ -49,6 +49,8 @@ def build_run_name(args): val_args = args.get("validation", {}) model_name = meta_args.get("model_name", "model") + representation_type = meta_args.get("representation_type") + head_type = meta_args.get("head_type") patch_size = mask_args.get("patch_size", meta_args.get("patch_size")) crop_size = data_args.get("crop_size", meta_args.get("crop_size")) batch_size = data_args.get("batch_size") @@ -67,6 +69,8 @@ def build_run_name(args): add("p", patch_size) add("c", crop_size) add("bs", batch_size) + add("rep", representation_type) + add("head", head_type) parts.append(str(optimizer).lower()) add("lr", lr) add("wd", weight_decay)