from __future__ import annotations """Unified experiment tracking: W&B + TensorBoard + CSV. Auto-detects available backends. Falls back gracefully if wandb/tensorboard are not installed. Usage: tracker = ExperimentTracker(output_dir, config_dict, use_wandb=True, use_tb=True) tracker.log_train(epoch, {"loss": 0.5, "lr": 1e-4}) tracker.log_val(epoch, {"r@1_q2g": 0.3}) tracker.log_gradients(epoch, grad_norms_dict) tracker.log_image(epoch, "gradcam/drone", image_tensor) tracker.close() """ import logging from pathlib import Path from typing import Any import torch LOGGER = logging.getLogger("caption_test.trackers") def _try_import_wandb(): try: import wandb return wandb except ImportError: return None def _try_import_tb(): try: from torch.utils.tensorboard import SummaryWriter return SummaryWriter except ImportError: return None class ExperimentTracker: """Unified tracker dispatching to W&B, TensorBoard, and CSV. Args: output_dir: Base output directory. config: Dict of hyperparameters to log. use_wandb: Enable Weights & Biases tracking. use_tb: Enable TensorBoard tracking. wandb_project: W&B project name. wandb_run_name: W&B run name (auto-generated if None). wandb_entity: W&B entity (team/user). """ def __init__( self, output_dir: str | Path, config: dict[str, Any] | None = None, use_wandb: bool = False, use_tb: bool = True, wandb_project: str = "caption-test-gtauav", wandb_run_name: str | None = None, wandb_entity: str | None = None, ) -> None: self.output_dir = Path(output_dir) self._wandb_run = None self._tb_writer = None # W&B init. if use_wandb: wandb = _try_import_wandb() if wandb is not None: self._wandb_run = wandb.init( project=wandb_project, name=wandb_run_name, entity=wandb_entity, config=config or {}, dir=str(self.output_dir), reinit=True, ) LOGGER.info("W&B initialized: %s", self._wandb_run.url) else: LOGGER.warning("wandb not installed, skipping W&B tracking") # TensorBoard init. if use_tb: SummaryWriter = _try_import_tb() if SummaryWriter is not None: tb_dir = self.output_dir / "tb_logs" tb_dir.mkdir(parents=True, exist_ok=True) self._tb_writer = SummaryWriter(log_dir=str(tb_dir)) LOGGER.info("TensorBoard initialized: %s", tb_dir) else: LOGGER.warning("tensorboard not installed, skipping TB tracking") @property def has_wandb(self) -> bool: return self._wandb_run is not None @property def has_tb(self) -> bool: return self._tb_writer is not None def log_train(self, epoch: int, metrics: dict[str, float], step: int | None = None) -> None: """Log training metrics for an epoch.""" if self._wandb_run is not None: self._wandb_run.log( {f"train/{k}": v for k, v in metrics.items()}, step=step or epoch, ) if self._tb_writer is not None: for k, v in metrics.items(): self._tb_writer.add_scalar(f"train/{k}", v, global_step=step or epoch) def log_val(self, epoch: int, metrics: dict[str, float], step: int | None = None) -> None: """Log validation metrics.""" if self._wandb_run is not None: self._wandb_run.log( {f"val/{k}": v for k, v in metrics.items()}, step=step or epoch, ) if self._tb_writer is not None: for k, v in metrics.items(): self._tb_writer.add_scalar(f"val/{k}", v, global_step=step or epoch) def log_gradients(self, epoch: int, grad_norms: dict[str, float], step: int | None = None) -> None: """Log gradient norms per parameter group.""" if self._wandb_run is not None: self._wandb_run.log( {f"gradients/{k}": v for k, v in grad_norms.items()}, step=step or epoch, ) if self._tb_writer is not None: for k, v in grad_norms.items(): self._tb_writer.add_scalar(f"gradients/{k}", v, global_step=step or epoch) def log_scalar(self, tag: str, value: float, step: int) -> None: """Log a single scalar.""" if self._wandb_run is not None: self._wandb_run.log({tag: value}, step=step) if self._tb_writer is not None: self._tb_writer.add_scalar(tag, value, global_step=step) def log_image(self, tag: str, image: Any, step: int, caption: str | None = None) -> None: """Log an image (numpy HWC or torch CHW). Args: tag: Image tag/name. image: numpy array [H,W,C] or torch tensor [C,H,W]. step: Global step. caption: Optional caption for W&B. """ if self._wandb_run is not None: wandb = _try_import_wandb() if isinstance(image, torch.Tensor): image_np = image.detach().cpu().permute(1, 2, 0).numpy() else: image_np = image self._wandb_run.log( {tag: wandb.Image(image_np, caption=caption)}, step=step, ) if self._tb_writer is not None: if isinstance(image, torch.Tensor): self._tb_writer.add_image(tag, image.detach().cpu(), global_step=step) else: self._tb_writer.add_image(tag, image, global_step=step, dataformats="HWC") def log_histogram(self, tag: str, values: torch.Tensor, step: int) -> None: """Log a histogram of values (weights, activations, etc.).""" if self._wandb_run is not None: wandb = _try_import_wandb() self._wandb_run.log( {tag: wandb.Histogram(values.detach().cpu().numpy())}, step=step, ) if self._tb_writer is not None: self._tb_writer.add_histogram(tag, values.detach().cpu(), global_step=step) def log_model_graph(self, model: torch.nn.Module, input_example: Any = None) -> None: """Log model graph to TensorBoard (if available).""" if self._tb_writer is not None and input_example is not None: try: self._tb_writer.add_graph(model, input_example) except Exception as e: LOGGER.warning("Failed to log model graph: %s", e) def watch_model(self, model: torch.nn.Module, log_freq: int = 100) -> None: """Enable W&B gradient/weight watching.""" if self._wandb_run is not None: wandb = _try_import_wandb() wandb.watch(model, log="all", log_freq=log_freq) def log_summary(self, summary: dict[str, Any]) -> None: """Log final summary metrics (best R@1, etc.).""" if self._wandb_run is not None: for k, v in summary.items(): self._wandb_run.summary[k] = v def close(self) -> None: """Flush and close all backends.""" if self._tb_writer is not None: self._tb_writer.flush() self._tb_writer.close() if self._wandb_run is not None: self._wandb_run.finish()