Improve prediction naming for experiment logging

Author: Donglai Wei

connectomics/decoding/tuning/optuna_tuner.py

@@ -89,8 +89,27 @@ def _resolve_tuning_prediction_files(
     return existing_files, [str(path) for path in expected_files]
 
 
+def _print_best_params_yaml(best_params_file: Path) -> None:
+    """Print the current best-params YAML to stdout for interactive tune runs."""
+    try:
+        best_params = OmegaConf.load(best_params_file)
+        yaml_text = OmegaConf.to_yaml(best_params).rstrip()
+    except Exception:
+        yaml_text = best_params_file.read_text().rstrip()
+
+    print("\n" + "=" * 80)
+    print(f"BEST PARAMETERS | {best_params_file}")
+    print("=" * 80)
+    if yaml_text:
+        print(yaml_text)
+    else:
+        print("[empty]")
+
+
 @contextmanager
-def _temporary_tuning_inference_overrides(*cfg_objects: Any):
+def _temporary_tuning_inference_overrides(
+    *cfg_objects: Any, checkpoint_path: str | None = None
+):
     """Force the pre-Optuna inference pass to cache raw predictions only."""
     inference_cfgs = []
     seen_inference_cfgs: set[int] = set()
@@ -109,7 +128,11 @@ def _temporary_tuning_inference_overrides(*cfg_objects: Any):
     if not inference_cfgs:
         raise ValueError("Missing runtime cfg.inference configuration required for tuning")
 
-    suffix = tta_cache_suffix(primary_cfg) if primary_cfg is not None else "_tta_x1_prediction.h5"
+    suffix = (
+        tta_cache_suffix(primary_cfg, checkpoint_path=checkpoint_path)
+        if primary_cfg is not None
+        else "_tta_x1_prediction.h5"
+    )
 
     backups = []
     for inference_cfg in inference_cfgs:
@@ -1286,6 +1309,7 @@ def run_tuning(model, trainer, cfg, checkpoint_path=None):
             "Delete this file to re-run tuning.",
             best_params_file,
         )
+        _print_best_params_yaml(best_params_file)
         return
 
     logger.info("STARTING PARAMETER TUNING | Output directory: %s", output_dir)
@@ -1297,7 +1321,8 @@ def run_tuning(model, trainer, cfg, checkpoint_path=None):
     logger.info("[1/4] Running inference on tuning dataset...")
 
     tune_data = cfg.data
-    cache_suffix = tta_cache_suffix(cfg)
+    prediction_checkpoint_path = getattr(model, "_prediction_checkpoint_path", None) or checkpoint_path
+    cache_suffix = tta_cache_suffix(cfg, checkpoint_path=prediction_checkpoint_path)
 
     output_pred_dir = cfg.inference.save_prediction.output_path
     predictions_dir = Path(output_pred_dir)
@@ -1329,7 +1354,11 @@ def run_tuning(model, trainer, cfg, checkpoint_path=None):
         # Run test to populate/load raw prediction caches only. Optuna applies its own
         # decoding sweep afterward, so the tune inference pass must not decode with the
         # default config first.
-        with _temporary_tuning_inference_overrides(cfg, getattr(model, "cfg", None)) as cache_suffix:
+        with _temporary_tuning_inference_overrides(
+            cfg,
+            getattr(model, "cfg", None),
+            checkpoint_path=prediction_checkpoint_path,
+        ) as cache_suffix:
             model._tune_mode = True
             try:
                 results = trainer.test(model, datamodule=datamodule, ckpt_path=checkpoint_path)
@@ -1456,6 +1485,8 @@ def run_tuning(model, trainer, cfg, checkpoint_path=None):
         study.best_value,
         study.best_params,
     )
+    if best_params_file.exists():
+        _print_best_params_yaml(best_params_file)
 
 
 def load_and_apply_best_params(cfg):

connectomics/training/lightning/__init__.py

@@ -19,12 +19,14 @@
 from .utils import (
     compute_tta_passes,
     extract_best_score_from_checkpoint,
+    final_prediction_output_tag,
     is_tta_cache_suffix,
     parse_args,
     resolve_prediction_cache_suffix,
     setup_config,
     setup_seed_everything,
     tta_cache_suffix,
+    tta_cache_suffix_candidates,
 )
 
 __all__ = [
@@ -46,7 +48,9 @@
     "setup_config",
     "extract_best_score_from_checkpoint",
     "compute_tta_passes",
+    "final_prediction_output_tag",
     "tta_cache_suffix",
+    "tta_cache_suffix_candidates",
     "resolve_prediction_cache_suffix",
     "is_tta_cache_suffix",
 ]

connectomics/training/lightning/model.py

@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import logging
+import os
 import warnings
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
@@ -44,7 +45,14 @@
 from ...models import build_model
 from ...models.loss import create_loss, get_loss_metadata_for_module
 from ..debugging import DebugManager
-from .utils import is_tta_cache_suffix, resolve_prediction_cache_suffix, tta_cache_suffix
+from .utils import (
+    final_prediction_output_tag,
+    format_checkpoint_name_tag,
+    is_tta_cache_suffix,
+    resolve_prediction_cache_suffix,
+    tta_cache_suffix,
+    tta_cache_suffix_candidates,
+)
 
 # Import training/inference components
 from ..loss import LossOrchestrator, build_loss_weighter, infer_num_loss_tasks_from_config
@@ -411,11 +419,36 @@ def _resolve_test_output_config(
         mode = "test"
         save_pred_cfg = self._get_runtime_inference_config().save_prediction
         output_dir_value = getattr(save_pred_cfg, "output_path", None)
-        cache_suffix = resolve_prediction_cache_suffix(self.cfg, mode=mode)
+        cache_suffix = resolve_prediction_cache_suffix(
+            self.cfg,
+            mode=mode,
+            checkpoint_path=self._get_prediction_checkpoint_path(),
+        )
 
         filenames = resolve_output_filenames(self.cfg, batch, global_step=self.global_step)
         return mode, output_dir_value, cache_suffix, filenames
 
+    def _get_prediction_checkpoint_path(self) -> str:
+        """Return the checkpoint/weights path whose stem should tag prediction caches."""
+        explicit_path = getattr(self, "_prediction_checkpoint_path", None)
+        if explicit_path is not None:
+            path_value = str(explicit_path).strip()
+            if path_value:
+                return path_value
+
+        trainer = getattr(self, "_trainer", None)
+        trainer_ckpt_path = getattr(trainer, "ckpt_path", None) if trainer is not None else None
+        if trainer_ckpt_path is not None:
+            path_value = str(trainer_ckpt_path).strip()
+            if path_value:
+                return path_value
+
+        external_weights_path = getattr(getattr(self.cfg, "model", None), "external_weights_path", None)
+        if isinstance(external_weights_path, str) and external_weights_path.strip():
+            return external_weights_path.strip()
+
+        return ""
+
     def _resolve_tta_result_path_override(self) -> str:
         """Return explicit intermediate prediction file from inference.tta_result_path."""
         inference_cfg = self._get_runtime_inference_config()
@@ -434,7 +467,7 @@ def _load_cached_predictions(
             if not pred_file.is_absolute():
                 pred_file = Path.cwd() / pred_file
 
-            if pred_file.exists():
+            if os.path.exists(pred_file):
                 try:
                     logger.info(f"Using explicit inference.tta_result_path file: {pred_file}")
                     pred = read_volume(str(pred_file), dataset="main")
@@ -446,7 +479,14 @@ def _load_cached_predictions(
                             f"{len(filenames)} filenames; decoding will use the explicit file only."
                         )
                     # Treat explicit file as intermediate prediction so decoding still runs.
-                    return pred, True, tta_cache_suffix(self.cfg)
+                    return (
+                        pred,
+                        True,
+                        tta_cache_suffix(
+                            self.cfg,
+                            checkpoint_path=self._get_prediction_checkpoint_path(),
+                        ),
+                    )
                 except Exception as e:
                     logger.warning(
                         f"Failed to load explicit inference.tta_result_path file {pred_file}: {e}. "
@@ -462,14 +502,22 @@ def _load_cached_predictions(
             return None, False, cache_suffix
 
         output_dir = Path(output_dir_value)
+        checkpoint_tag = format_checkpoint_name_tag(self._get_prediction_checkpoint_path())
 
         # Build ordered list of suffixes to try: final prediction first, then
         # intermediate TTA, then glob fallback.
         suffixes_to_try: list[str] = []
         if is_tta_cache_suffix(cache_suffix):
             # Prefer the final decoded file (e.g. _x16_prediction.h5) over
             # the intermediate TTA file (e.g. _tta_x16_prediction.h5).
-            final_suffix = cache_suffix.replace("_tta_x", "_x")
+            final_suffix = (
+                "_"
+                + final_prediction_output_tag(
+                    self.cfg,
+                    checkpoint_path=self._get_prediction_checkpoint_path(),
+                )
+                + ".h5"
+            )
             suffixes_to_try.append(final_suffix)
         suffixes_to_try.append(cache_suffix)
 
@@ -478,7 +526,7 @@ def _load_cached_predictions(
             all_exist = True
             for filename in filenames:
                 pred_file = output_dir / f"{filename}{try_suffix}"
-                if pred_file.exists():
+                if os.path.exists(pred_file):
                     try:
                         pred = read_volume(str(pred_file), dataset="main")
                         existing_predictions.append(pred)
@@ -507,21 +555,43 @@ def _load_cached_predictions(
                     )
                 return predictions_np, True, try_suffix
 
-        # Glob fallback: look for any TTA intermediate file.
+        # Targeted fallback: look for the exact TTA intermediate cache suffix
+        # matching the current config rather than any arbitrary TTA file.
         if mode == "test" and not is_tta_cache_suffix(cache_suffix):
-            for filename in filenames:
-                tta_matches = sorted(output_dir.glob(f"{filename}_tta_x*_prediction.h5"))
-                if tta_matches:
-                    pred_file = tta_matches[-1]
-                    loaded_suffix = pred_file.name[len(filename):]
+            fallback_suffixes = tta_cache_suffix_candidates(
+                self.cfg,
+                checkpoint_path=self._get_prediction_checkpoint_path(),
+            )
+            for try_suffix in fallback_suffixes:
+                existing_predictions = []
+                all_exist = True
+                for filename in filenames:
+                    pred_file = output_dir / f"{filename}{try_suffix}"
+                    if not os.path.exists(pred_file):
+                        all_exist = False
+                        break
                     try:
                         pred = read_volume(str(pred_file), dataset="main")
-                        if pred.ndim < 4:
-                            pred = pred[np.newaxis, ...]
-                        logger.info("Loaded fallback TTA prediction: %s", pred_file.name)
-                        return pred, True, loaded_suffix
+                        existing_predictions.append(pred)
                     except Exception as e:
                         logger.warning(f"Failed to load {pred_file}: {e}")
+                        all_exist = False
+                        break
+                if all_exist and len(existing_predictions) == len(filenames):
+                    logger.info(
+                        "Loaded fallback TTA prediction(s) using exact suffix %s",
+                        try_suffix,
+                    )
+                    if len(existing_predictions) == 1:
+                        predictions_np = existing_predictions[0]
+                        if predictions_np.ndim < 4:
+                            predictions_np = predictions_np[np.newaxis, ...]
+                    else:
+                        predictions_np = np.stack(
+                            [p[np.newaxis, ...] if p.ndim < 4 else p for p in existing_predictions],
+                            axis=0,
+                        )
+                    return predictions_np, True, try_suffix
 
         return None, False, cache_suffix
 
@@ -554,9 +624,10 @@ def _save_metrics_to_file(self, metrics_dict: Dict[str, Any]):
         # Create filename with volume name and TTA pass tag
         volume_name = metrics_dict.get("volume_name", "unknown")
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        cache_suffix = resolve_prediction_cache_suffix(self.cfg, mode=mode)
-        # Extract tag like "tta_x16" or "x1" from suffix "_tta_x16_prediction.h5"
-        tag = cache_suffix.lstrip("_").replace("_prediction.h5", "")
+        tag = final_prediction_output_tag(
+            self.cfg,
+            checkpoint_path=self._get_prediction_checkpoint_path(),
+        )
         metrics_file = output_dir / f"evaluation_metrics_{volume_name}_{tag}.txt"
 
         # Write metrics to file
@@ -663,11 +734,15 @@ def _log_decode_experiment(
             decode_params["decoder"] = step.name
             decode_params.update(step.kwargs)
 
+        input_tta_prediction_name = (
+            f"{volume_name}{tta_cache_suffix(self.cfg, checkpoint_path=self._get_prediction_checkpoint_path())}"
+        )
+
         # Columns: timestamp, volume, decoder params..., metrics...
         # Use a fixed column order for readability
         param_keys = [
             "decoder", "thresholds", "merge_function", "aff_threshold",
-            "channel_order", "dust_merge_size", "dust_merge_affinity",
+            "channel_order", "dust_merge", "dust_merge_size", "dust_merge_affinity",
             "dust_remove_size",
         ]
         metric_keys = [
@@ -676,8 +751,8 @@ def _log_decode_experiment(
             "instance_f1_detail",
         ]
 
-        header_cols = ["timestamp", "volume"] + param_keys + metric_keys
-        row_vals = [timestamp, volume_name]
+        header_cols = ["timestamp", "volume", "input_tta_prediction_name"] + param_keys + metric_keys
+        row_vals = [timestamp, volume_name, input_tta_prediction_name]
         for k in param_keys:
             row_vals.append(str(decode_params.get(k, "")))
         for k in metric_keys:

connectomics/training/lightning/test_pipeline.py

@@ -28,7 +28,14 @@
 )
 from ...metrics.metrics_seg import AdaptedRandError
 from ...metrics.segmentation_numpy import instance_matching, instance_matching_simple, voi
-from .utils import compute_tta_passes, is_tta_cache_suffix
+from .utils import (
+    compute_tta_passes,
+    final_prediction_output_tag,
+    format_checkpoint_name_tag,
+    format_decode_tag,
+    format_select_channel_tag,
+    is_tta_cache_suffix,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -790,12 +797,14 @@ def _process_decoding_postprocessing(
     if save_final_predictions:
         logger.info("[STAGE: Saving Final Predictions]")
         save_start = time.time()
-        final_tta_passes = compute_tta_passes(module.cfg)
         write_outputs(
             module.cfg,
             postprocessed_predictions,
             filenames,
-            suffix=f"x{final_tta_passes}_prediction",
+            suffix=final_prediction_output_tag(
+                module.cfg,
+                checkpoint_path=module._get_prediction_checkpoint_path(),
+            ),
             mode=mode,
             batch_meta=batch_meta,
         )
@@ -1058,12 +1067,14 @@ def run_test_step(module, batch: Dict[str, torch.Tensor], batch_idx: int) -> STE
         logger.info("[STAGE: Saving Intermediate Predictions]")
         save_start = time.time()
         tta_passes = compute_tta_passes(module.cfg)
+        ch_tag = format_select_channel_tag(module.cfg)
+        checkpoint_tag = format_checkpoint_name_tag(module._get_prediction_checkpoint_path())
         predictions_to_save = apply_save_prediction_transform(module.cfg, predictions_np)
         write_outputs(
             module.cfg,
             predictions_to_save,
             filenames,
-            suffix=f"tta_x{tta_passes}_prediction",
+            suffix=f"tta_x{tta_passes}{ch_tag}{checkpoint_tag}_prediction",
             mode=mode,
             batch_meta=batch.get("image_meta_dict"),
         )