From 17335bcc3dd8710b82c7efd806b57a7b38bc682c Mon Sep 17 00:00:00 2001 From: Lufang Chen <64068400+vincentccc@users.noreply.github.com> Date: Sat, 27 Sep 2025 12:51:25 +0800 Subject: [PATCH] Update train test split (#41) * update * lint --- wall_x/data/load_lerobot_dataset.py | 94 +++++++++++++++++++++------ wall_x/trainer/qwen_vl_act_trainer.py | 2 + 2 files changed, 76 insertions(+), 20 deletions(-) diff --git a/wall_x/data/load_lerobot_dataset.py b/wall_x/data/load_lerobot_dataset.py index 56ad613..f5ce2fb 100644 --- a/wall_x/data/load_lerobot_dataset.py +++ b/wall_x/data/load_lerobot_dataset.py @@ -2,9 +2,10 @@ LeRobot Dataset Loader - Distributed Version """ +import numpy as np import torch -from torch.utils.data import DistributedSampler -from lerobot.datasets.lerobot_dataset import LeRobotDataset +from torch.utils.data import DistributedSampler, random_split +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata from typing import Protocol, SupportsIndex, TypeVar from qwen_vl_utils.vision_process import smart_resize from wall_x.data.config import X2RDataProcessingConfig @@ -40,8 +41,29 @@ class Dataset(Protocol[T_co]): class PreprocessedDataset(Dataset[T_co]): - def __init__(self, dataset, config, dataload_config, seed=42, rank=0, world_size=1): - self._dataset = dataset + def __init__( + self, + dataset, + config, + dataload_config, + seed=42, + rank=0, + world_size=1, + test_only=False, + ): + self.hf_dataset = dataset + + if test_only: + self._dataset = dataset + else: + self._dataset = None + self.train_dataset, self.val_dataset = random_split( + dataset, + [0.95, 0.05], + torch.Generator().manual_seed(seed) if seed is not None else None, + ) + self._train() + self.seed = seed self.rank = rank self.world_size = world_size @@ -60,11 +82,11 @@ class PreprocessedDataset(Dataset[T_co]): priority_order=self.dataload_config.get("priority_order", None), ) - self._cam_key_mapping = CAMERA_KEY_MAPPINGS[self._dataset.meta.repo_id] + self._cam_key_mapping = CAMERA_KEY_MAPPINGS[self.hf_dataset.meta.repo_id] def _vision_preprocess(self, frames): processed_frames = [] - for key in self._dataset.meta.camera_keys: + for key in self.hf_dataset.meta.camera_keys: from PIL import Image current_obs = frames[key].clone().permute(1, 2, 0) @@ -131,6 +153,12 @@ class PreprocessedDataset(Dataset[T_co]): def __len__(self) -> int: return len(self._dataset) + def _eval(self): + self._dataset = self.val_dataset + + def _train(self): + self._dataset = self.train_dataset + def get_train_dataloader(self): """ Get distributed training dataloader @@ -140,6 +168,7 @@ class PreprocessedDataset(Dataset[T_co]): world_size: Total number of processes seed: Random seed for reproducibility """ + self._train() batch_size = self.config.get("batch_size_per_gpu", 8) num_workers = self.config.get("num_workers", 4) @@ -160,7 +189,7 @@ class PreprocessedDataset(Dataset[T_co]): sampler=sampler, # Use distributed sampler instead of shuffle=True num_workers=num_workers, collate_fn=DataCollator( - self.config, self.dataload_config, self._dataset.meta.stats + self.config, self.dataload_config, self.hf_dataset.meta.stats ), pin_memory=True, # Enable for GPU training persistent_workers=num_workers > 0, # Only if num_workers > 0 @@ -174,6 +203,7 @@ class PreprocessedDataset(Dataset[T_co]): """ Get distributed evaluation dataloader (no shuffling for consistent evaluation) """ + self._eval() batch_size = self.config.get( "eval_batch_size_per_gpu", self.config.get("batch_size_per_gpu", 8) @@ -195,7 +225,7 @@ class PreprocessedDataset(Dataset[T_co]): sampler=sampler, num_workers=num_workers, collate_fn=DataCollator( - self.config, self.dataload_config, self._dataset.meta.stats + self.config, self.dataload_config, self.hf_dataset.meta.stats ), pin_memory=True, persistent_workers=num_workers > 0, @@ -410,9 +440,15 @@ def load_lerobot_data( # Set seed for reproducibility torch.manual_seed(seed) - dataset_fps = 50 dataload_config = get_data_configs(config["data"]) + # repo_id = "lerobot/aloha_mobile_cabinet" + repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet") + root = lerobot_config.get("root", None) + meta_info = LeRobotDatasetMetadata(repo_id) + dataset_fps = meta_info.fps + episodes_num = meta_info.total_episodes + delta_timestamps = { # action chunk "action": [ @@ -421,21 +457,33 @@ def load_lerobot_data( ], } batch_size = config.get("batch_size_per_gpu", 8) + episodes = np.arange(episodes_num).tolist() - # repo_id = "lerobot/aloha_mobile_cabinet" - repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet") - root = lerobot_config.get("root", None) - dataset = LeRobotDataset( - repo_id, root=root, delta_timestamps=delta_timestamps, video_backend="pyav" + train_test_split = dataload_config.get("train_test_split", 0.95) + train_episodes = episodes[: int(episodes_num * train_test_split)] + test_episodes = episodes[int(episodes_num * train_test_split) :] + + train_dataset = LeRobotDataset( + repo_id, + root=root, + episodes=train_episodes, + delta_timestamps=delta_timestamps, + video_backend="pyav", ) if rank == 0: - print(f"Selected episodes: {dataset.episodes}") - print(f"Number of episodes selected: {dataset.num_episodes}") - print(f"Number of frames selected: {dataset.num_frames}") + print(f"Selected train episodes: {train_dataset.episodes}") + print(f"Number of train episodes selected: {train_dataset.num_episodes}") + print(f"Number of train frames selected: {train_dataset.num_frames}") + print(f"Selected test episodes: {test_episodes}") dataset = PreprocessedDataset( - dataset, config, dataload_config, seed=seed, rank=rank, world_size=world_size + train_dataset, + config, + dataload_config, + seed=seed, + rank=rank, + world_size=world_size, ) # Calculate samples per process @@ -516,7 +564,13 @@ def get_data_configs(config): class TestDataset(PreprocessedDataset): def __init__(self, dataset, config, dataload_config, seed=42): super().__init__( - dataset, config, dataload_config, seed=seed, rank=0, world_size=1 + dataset, + config, + dataload_config, + seed=seed, + rank=0, + world_size=1, + test_only=True, ) def get_dataloader(self): @@ -528,7 +582,7 @@ class TestDataset(PreprocessedDataset): self, batch_size=1, collate_fn=DataCollator( - self.config, self.dataload_config, self._dataset.meta.stats + self.config, self.dataload_config, self.hf_dataset.meta.stats ), ) diff --git a/wall_x/trainer/qwen_vl_act_trainer.py b/wall_x/trainer/qwen_vl_act_trainer.py index 92e4a2a..18dd181 100644 --- a/wall_x/trainer/qwen_vl_act_trainer.py +++ b/wall_x/trainer/qwen_vl_act_trainer.py @@ -253,6 +253,7 @@ class QwenVlAct_Trainer: # Initialize training dataloader for current epoch if isinstance(self.dataset, PreprocessedDataset): if getattr(self, "train_dataloader", None) is not None: + self.dataset._train() self.train_sampler.set_epoch(epoch) else: self.train_dataloader, self.train_sampler = ( @@ -463,6 +464,7 @@ class QwenVlAct_Trainer: """ # Initialize validation dataloader if getattr(self, "val_dataloader", None) is not None: + self.dataset._eval() self.val_sampler.set_epoch(0) else: self.val_dataloader, self.val_sampler = self.dataset.get_val_dataloader()