fuze_task/vendor_reference/caption_test/src/training/trackers.py

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()