mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 02:06:15 +00:00
refactor(train): remove wandb example tables
This commit is contained in:
@@ -19,8 +19,6 @@ import re
|
|||||||
from glob import glob
|
from glob import glob
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
import torch
|
|
||||||
from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
|
from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
|
||||||
from termcolor import colored
|
from termcolor import colored
|
||||||
|
|
||||||
@@ -209,86 +207,3 @@ class WandBLogger:
|
|||||||
|
|
||||||
wandb_video = self._wandb.Video(video_path, fps=self.env_fps, format="mp4")
|
wandb_video = self._wandb.Video(video_path, fps=self.env_fps, format="mp4")
|
||||||
self._wandb.log({f"{mode}/video": wandb_video}, step=step)
|
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)
|
|
||||||
|
|||||||
@@ -75,9 +75,6 @@ class WandBConfig:
|
|||||||
run_id: str | None = None
|
run_id: str | None = None
|
||||||
mode: str | None = None # Allowed values: 'online', 'offline' 'disabled'. Defaults to 'online'
|
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.
|
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
|
@dataclass
|
||||||
|
|||||||
@@ -730,23 +730,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
|||||||
if wandb_logger:
|
if wandb_logger:
|
||||||
wandb_logger.log_dict({"eval_loss": eval_loss}, step=step, mode="eval")
|
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:
|
if cfg.save_checkpoint and is_saving_step:
|
||||||
# Under FSDP, gathering the full model + optimizer state dicts is a cross-rank collective,
|
# 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 /
|
# so all ranks must participate; rank 0 then writes the materialized dicts. For DDP /
|
||||||
|
|||||||
Reference in New Issue
Block a user