refactor(train): remove wandb example tables

This commit is contained in:
Pepijn
2026-07-15 14:05:50 +02:00
parent 07e75d94be
commit 0a7b21cdd0
3 changed files with 0 additions and 105 deletions
-85
View File
@@ -19,8 +19,6 @@ import re
from glob import glob
from pathlib import Path
import numpy as np
import torch
from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
from termcolor import colored
@@ -209,86 +207,3 @@ class WandBLogger:
wandb_video = self._wandb.Video(video_path, fps=self.env_fps, format="mp4")
self._wandb.log({f"{mode}/video": wandb_video}, step=step)
def log_training_examples(
self,
batch: dict,
step: int,
*,
camera_keys: list[str],
n_samples: int = 4,
mode: str = "train",
) -> None:
"""Log a small W&B table with sampled images/text and action endpoints."""
if mode not in {"train", "eval"}:
raise ValueError(mode)
# Batch size — first tensor-like value wins.
bsz = next(
(int(v.shape[0]) for v in batch.values() if hasattr(v, "shape") and v.ndim > 0),
None,
)
if not bsz:
return
n = min(int(n_samples), bsz)
present_cameras = [c for c in camera_keys if c in batch]
text_keys = [k for k in ("task", "subtask", "memory", "instruction") if k in batch]
columns = ["sample"]
columns.extend(c.removeprefix("observation.images.") or c for c in present_cameras)
columns.extend(text_keys)
columns += ["gt_action_first", "gt_action_last"]
table = self._wandb.Table(columns=columns)
def _to_uint8_hwc(t: torch.Tensor) -> np.ndarray:
if t.ndim == 4:
t = t[0]
if t.ndim == 3 and t.shape[0] in (1, 3, 4) and t.shape[-1] not in (1, 3, 4):
t = t.permute(1, 2, 0)
arr = t.detach().cpu().float().numpy()
if arr.size and float(arr.max()) <= 1.5:
arr = arr * 255.0
return np.clip(arr, 0, 255).astype(np.uint8)
def _action_endpoints(a: torch.Tensor) -> tuple[str, str]:
arr = a.detach().cpu().float().numpy()
if arr.ndim == 2:
return str(np.round(arr[0], 3).tolist()), str(np.round(arr[-1], 3).tolist())
if arr.ndim == 1:
rounded = np.round(arr, 3).tolist()
return str(rounded), str(rounded)
text = str(arr.tolist())
return text, text
for i in range(n):
row: list = [i]
for cam in present_cameras:
try:
row.append(self._wandb.Image(_to_uint8_hwc(batch[cam][i])))
except Exception as exc: # noqa: BLE001
logging.warning(
"log_training_examples: camera %s sample %d failed (%s)",
cam,
i,
exc,
)
row.append(None)
for tk in text_keys:
v = batch[tk]
if isinstance(v, list | tuple):
row.append(str(v[i]) if i < len(v) else "")
else:
row.append(str(v))
action = batch.get("action")
if isinstance(action, torch.Tensor) and action.ndim >= 1:
first, last = _action_endpoints(action[i])
row.append(first)
row.append(last)
else:
row.append("")
row.append("")
table.add_data(*row)
self._wandb.log({f"{mode}/examples": table}, step=step)
-3
View File
@@ -75,9 +75,6 @@ class WandBConfig:
run_id: str | None = None
mode: str | None = None # Allowed values: 'online', 'offline' 'disabled'. Defaults to 'online'
add_tags: bool = True # If True, save configuration as tags in the WandB run.
# Periodic W&B table with sampled images/text and action endpoints. Set to 0 to disable.
log_examples_freq: int = 5000
log_examples_n: int = 4
@dataclass
-17
View File
@@ -730,23 +730,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if wandb_logger:
wandb_logger.log_dict({"eval_loss": eval_loss}, step=step, mode="eval")
# Periodic W&B example table (camera images + text fields + action endpoints).
if (
wandb_logger is not None
and cfg.wandb.log_examples_freq > 0
and step % cfg.wandb.log_examples_freq == 0
and is_main_process
):
try:
wandb_logger.log_training_examples(
batch=batch,
step=step,
camera_keys=list(dataset.meta.camera_keys),
n_samples=cfg.wandb.log_examples_n,
)
except Exception as exc: # noqa: BLE001
logging.warning("wandb log_training_examples failed: %s", exc)
if cfg.save_checkpoint and is_saving_step:
# Under FSDP, gathering the full model + optimizer state dicts is a cross-rank collective,
# so all ranks must participate; rank 0 then writes the materialized dicts. For DDP /