数据集检查

import os
from typing import Any

import hydra
import torch
from hydra.utils import instantiate
from omegaconf import DictConfig, OmegaConf
from torch.utils.data import DataLoader

from lightwam.utils.config_resolvers import register_default_resolvers

register_default_resolvers()

def _summarize_value(value: Any) -> str:
    if isinstance(value, torch.Tensor):
        return f"Tensor shape={tuple(value.shape)} dtype={value.dtype} device={value.device}"
    if isinstance(value, (str, bytes)):
        text = value.decode("utf-8", errors="replace") if isinstance(value, bytes) else value
        preview = text[:120].replace("\n", "\\n")
        suffix = "..." if len(text) > 120 else ""
        return f"{type(value).__name__} len={len(text)} preview='{preview}{suffix}'"
    if isinstance(value, (list, tuple)):
        return f"{type(value).__name__} len={len(value)}"
    if isinstance(value, dict):
        return f"dict keys={sorted(list(value.keys()))[:20]}"
    return f"{type(value).__name__} value={value}"


def _summarize_batch(batch: dict[str, Any], prefix: str = ""):
    keys = sorted(batch.keys())
    print(f"{prefix}keys={keys}")
    for key in keys:
        print(f"{prefix}{key}: {_summarize_value(batch[key])}")


@hydra.main(config_path="../configs", config_name="train", version_base="1.3")
def main(cfg: DictConfig):
    os.environ.setdefault("HYDRA_FULL_ERROR", "1")

    inspect_cfg = cfg.get("inspect", {})
    num_batches = int(inspect_cfg.get("num_batches", 2))
    max_items_per_batch = int(inspect_cfg.get("max_items_per_batch", 2))
    dataloader_workers = int(inspect_cfg.get("num_workers", cfg.get("num_workers", 0)))
    dataloader_batch_size = int(inspect_cfg.get("batch_size", 2))

    print("[inspect] resolved config (data.train)")
    train_cfg_plain = OmegaConf.to_container(cfg.data.train, resolve=True)
    print(OmegaConf.to_yaml(train_cfg_plain, resolve=True))

    train_ds = instantiate(cfg.data.train)
    print(f"[inspect] dataset={type(train_ds).__name__} len={len(train_ds)}")

    loader = DataLoader(
        train_ds,
        batch_size=dataloader_batch_size,
        shuffle=False,
        num_workers=dataloader_workers,
        pin_memory=torch.cuda.is_available(),
    )

    for batch_idx, batch in enumerate(loader):
        if batch_idx >= num_batches:
            break
        print(f"[inspect] batch_idx={batch_idx}")

        if isinstance(batch, dict):
            _summarize_batch(batch, prefix="  ")
            if "video" in batch:
                print("  [inspect] uses=sample['video']")
            elif "video_latents" in batch:
                print("  [inspect] uses=sample['video_latents']")
            else:
                print("  [inspect] uses=<missing video/video_latents>")

            if max_items_per_batch > 0 and "idx" in batch:
                try:
                    idx_val = batch["idx"]
                    idx_list = idx_val[:max_items_per_batch].tolist() if isinstance(idx_val, torch.Tensor) else idx_val
                    print(f"  [inspect] idx_sample={idx_list}")
                except Exception:
                    pass
        else:
            print(f"  [inspect] unexpected batch type: {type(batch)}")


if __name__ == "__main__":
    main()

运行命令

python scripts/inspect_dataset_loading.py \
  task=libero_uncond_2cam224_1e-4 \
  data.train.dataset_dirs="['./data/libero_mujoco3.3.2/libero_goal_no_noops_lerobot']" \
  data.train.text_embedding_cache_dir='./data/text_embeds_cache/libero' \
  data.train.use_latent_cache=true \
  data.train.latent_cache_dir='./data/latent_cache_Wan2.1-T2V-1.3B/libero_goal_2cam224' \
  +inspect.num_batches=2 \
  +inspect.batch_size=2 \
  +inspect.num_workers=4

[inspect] resolved config (data.train)
_target_: lightwam.datasets.lerobot.robot_video_dataset.RobotVideoDataset
dataset_dirs:
- ./data/libero_mujoco3.3.2/libero_goal_no_noops_lerobot
shape_meta:
  images:
  - key: image
    raw_shape:
    - 3
    - 512
    - 512
    shape:
    - 3
    - 224
    - 224
  - key: wrist_image
    raw_shape:
    - 3
    - 512
    - 512
    shape:
    - 3
    - 224
    - 224
  action:
  - key: default
    raw_shape: 7
    shape: 7
  state:
  - key: default
    raw_shape: 8
    shape: 8
num_frames: 33
global_sample_stride: 1
action_video_freq_ratio: 4
video_size:
- 224
- 448
camera_key: null
val_set_proportion: 0.0
is_training_set: true
pretrained_norm_stats: null
skip_padding_as_possible: false
concat_multi_camera: horizontal
use_latent_cache: true
latent_cache_dir: ./data/latent_cache_Wan2.1-T2V-1.3B/libero_goal_2cam224
processor:
  _target_: lightwam.datasets.lerobot.processors.lightwam_processor.LightWAMProcessor
  shape_meta:
    images:
    - key: image
      raw_shape:
      - 3
      - 512
      - 512
      shape:
      - 3
      - 224
      - 224
    - key: wrist_image
      raw_shape:
      - 3
      - 512
      - 512
      shape:
      - 3
      - 224
      - 224
    action:
    - key: default
      raw_shape: 7
      shape: 7
    state:
    - key: default
      raw_shape: 8
      shape: 8
  num_obs_steps: 33
  num_output_cameras: 2
  action_output_dim: 7
  proprio_output_dim: 8
  delta_action_dim_mask:
    default:
    - true
    - true
    - true
    - true
    - true
    - true
    - false
  action_state_transforms: null
  use_stepwise_action_norm: false
  norm_default_mode: min/max
  norm_exception_mode: null
  action_state_merger:
    _target_: lightwam.datasets.lerobot.transforms.action_state_merger.ConcatLeftAlign
  train_transforms:
  - _target_: lightwam.datasets.lerobot.transforms.image.ToTensor
  - _target_: torchvision.transforms.Resize
    size:
    - 224
    - 224
  val_transforms:
  - _target_: lightwam.datasets.lerobot.transforms.image.ToTensor
  - _target_: torchvision.transforms.Resize
    size:
    - 224
    - 224
text_embedding_cache_dir: ./data/text_embeds_cache/libero
context_len: 128

[2026-07-27 08:40:02,634][datasets][INFO] - PyTorch version 2.7.1+cu128 available.
Resolving data files: 100%|█████████████████████████████████████████████████████████████████████████████████| 433/433 [00:00<00:00, 477301.87it/s]
Downloading data: 100%|███████████████████████████████████████████████████████████████████████████████████| 433/433 [00:00<00:00, 33418.60files/s]
Generating train split: 52895 examples [00:00, 73563.37 examples/s] 
[2026-07-27 08:40:08,523][lightwam.datasets.lerobot.robot_video_dataset][INFO] - Calculating dataset stats for normalization...
Iterating dataset to get normalization: 100%|███████████████████████████████████████████████████████████████████| 433/433 [00:04<00:00, 96.65it/s]
[2026-07-27 08:40:13,074][lightwam.datasets.lerobot.robot_video_dataset][INFO] - Loaded indexed latent cache index: format=sharded_v1 shards=52 samples=52895
[2026-07-27 08:40:13,075][lightwam.datasets.lerobot.robot_video_dataset][INFO] - Using latent cache for RobotVideoDataset: /workspace/Light-WAM/data/latent_cache_Wan2.1-T2V-1.3B/libero_goal_2cam224
[inspect] dataset=RobotVideoDataset len=52895
[inspect] batch_idx=0
  keys=['action', 'action_is_pad', 'context', 'context_mask', 'idx', 'image_is_pad', 'prompt', 'proprio', 'proprio_is_pad', 'video_latents']
  action: Tensor shape=(2, 32, 7) dtype=torch.float32 device=cpu
  action_is_pad: Tensor shape=(2, 32) dtype=torch.bool device=cpu
  context: Tensor shape=(2, 128, 4096) dtype=torch.bfloat16 device=cpu
  context_mask: Tensor shape=(2, 128) dtype=torch.bool device=cpu
  idx: Tensor shape=(2,) dtype=torch.int64 device=cpu
  image_is_pad: Tensor shape=(2, 9) dtype=torch.bool device=cpu
  prompt: list len=2
  proprio: Tensor shape=(2, 32, 8) dtype=torch.float32 device=cpu
  proprio_is_pad: Tensor shape=(2, 33) dtype=torch.bool device=cpu
  video_latents: Tensor shape=(2, 16, 3, 28, 56) dtype=torch.bfloat16 device=cpu
  [inspect] uses=sample['video_latents']
  [inspect] idx_sample=[0, 1]
[inspect] batch_idx=1
  keys=['action', 'action_is_pad', 'context', 'context_mask', 'idx', 'image_is_pad', 'prompt', 'proprio', 'proprio_is_pad', 'video_latents']
  action: Tensor shape=(2, 32, 7) dtype=torch.float32 device=cpu
  action_is_pad: Tensor shape=(2, 32) dtype=torch.bool device=cpu
  context: Tensor shape=(2, 128, 4096) dtype=torch.bfloat16 device=cpu
  context_mask: Tensor shape=(2, 128) dtype=torch.bool device=cpu
  idx: Tensor shape=(2,) dtype=torch.int64 device=cpu
  image_is_pad: Tensor shape=(2, 9) dtype=torch.bool device=cpu
  prompt: list len=2
  proprio: Tensor shape=(2, 32, 8) dtype=torch.float32 device=cpu
  proprio_is_pad: Tensor shape=(2, 33) dtype=torch.bool device=cpu
  video_latents: Tensor shape=(2, 16, 3, 28, 56) dtype=torch.bfloat16 device=cpu
  [inspect] uses=sample['video_latents']
  [inspect] idx_sample=[2, 3]