From f1d2c9e6b76007c10ff10181f0ee8d59182e39cc Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Fri, 12 Dec 2025 05:48:04 +0000 Subject: [PATCH] Add RL dataset & dataloader --- fastvideo/dataset/__init__.py | 4 +- fastvideo/dataset/rl_prompt_dataset.py | 171 +++++++++++++++++++++++++ 2 files changed, 174 insertions(+), 1 deletion(-) create mode 100644 fastvideo/dataset/rl_prompt_dataset.py diff --git a/fastvideo/dataset/__init__.py b/fastvideo/dataset/__init__.py index 9c1e98473..c7dfc3ab0 100644 --- a/fastvideo/dataset/__init__.py +++ b/fastvideo/dataset/__init__.py @@ -8,6 +8,7 @@ from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset, from fastvideo.dataset.transform import (CenterCropResizeVideo, Normalize255, TemporalRandomCrop) from fastvideo.dataset.validation_dataset import ValidationDataset +from fastvideo.dataset.rl_prompt_dataset import build_rl_prompt_dataloader def getdataset(args) -> VideoCaptionMergedDataset: @@ -47,5 +48,6 @@ def gettextdataset(args) -> TextDataset: __all__ = [ "build_parquet_map_style_dataloader", "ValidationDataset", - "VideoCaptionMergedDataset", "TextDataset" + "VideoCaptionMergedDataset", "TextDataset", + "build_rl_prompt_dataloader" ] diff --git a/fastvideo/dataset/rl_prompt_dataset.py b/fastvideo/dataset/rl_prompt_dataset.py new file mode 100644 index 000000000..a67094f58 --- /dev/null +++ b/fastvideo/dataset/rl_prompt_dataset.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +import torch +from torch.utils.data import Dataset, DataLoader, Sampler +import json +import os + + +class TextPromptDataset(Dataset): + """Dataset for loading text prompts from a simple text file (one prompt per line).""" + + def __init__(self, dataset, split='train'): + self.file_path = os.path.join(dataset, f'{split}.txt') + with open(self.file_path, 'r') as f: + self.prompts = [line.strip() for line in f.readlines()] + + def __len__(self): + return len(self.prompts) + + def __getitem__(self, idx): + return {"prompt": self.prompts[idx], "metadata": {}} + + @staticmethod + def collate_fn(examples): + prompts = [example["prompt"] for example in examples] + metadatas = [example["metadata"] for example in examples] + return prompts, metadatas + + +class GenevalPromptDataset(Dataset): + """Dataset for loading prompts with metadata from JSONL files (e.g., GenEval format).""" + + def __init__(self, dataset, split='train'): + self.file_path = os.path.join(dataset, f'{split}_metadata.jsonl') + with open(self.file_path, 'r', encoding='utf-8') as f: + self.metadatas = [json.loads(line) for line in f] + self.prompts = [item['prompt'] for item in self.metadatas] + + def __len__(self): + return len(self.prompts) + + def __getitem__(self, idx): + return {"prompt": self.prompts[idx], "metadata": self.metadatas[idx]} + + @staticmethod + def collate_fn(examples): + prompts = [example["prompt"] for example in examples] + metadatas = [example["metadata"] for example in examples] + return prompts, metadatas + + +class KRepeatSampler(Sampler): + """Sampler that repeats each sample k times, ensuring synchronized random selection. For single-node training, set num_replicas=1 and rank=0.""" + + def __init__(self, dataset, batch_size, k, num_replicas, rank, seed=0): + self.dataset = dataset + self.batch_size = batch_size # Batch size per GPU/card + self.k = k # Number of repetitions per sample + self.num_replicas = num_replicas # Total number of GPUs/cards + self.rank = rank # Current GPU/card rank + self.seed = seed # Random seed for synchronization + + # Calculate the number of unique samples needed for each iteration + self.total_samples = self.num_replicas * self.batch_size + assert self.total_samples % self.k == 0, f"k can not div n*b, k{k}-num_replicas{num_replicas}-batch_size{batch_size}" + self.m = self.total_samples // self.k # different number of samples + self.step = 0 + + def __iter__(self): + while True: + # Generate a deterministic random sequence to ensure all cards are synchronized + g = torch.Generator() + g.manual_seed(self.seed + self.step) + + # Randomly select m unique samples + indices = torch.randperm(len(self.dataset), generator=g)[:self.m].tolist() + + # Repeat each sample k times to generate a total of n*b samples + repeated_indices = [idx for idx in indices for _ in range(self.k)] + + # Shuffle the order to ensure even distribution + shuffled_indices = torch.randperm(len(repeated_indices), generator=g).tolist() + shuffled_samples = [repeated_indices[i] for i in shuffled_indices] + + # Split samples among all cards + per_card_samples = [] + for i in range(self.num_replicas): + start = i * self.batch_size + end = start + self.batch_size + per_card_samples.append(shuffled_samples[start:end]) + + # Return the sample indices for the current card + yield per_card_samples[self.rank] + + def set_step(self, step): + """Used to synchronize the random state for different epochs.""" + self.step = step + + +def build_rl_prompt_dataloader( + dataset_path: str, + dataset_type: str = "text", + split: str = "train", + train_batch_size: int = 8, + test_batch_size: int = 8, + k: int = 1, + seed: int = 42, + train_num_workers: int = 1, + test_num_workers: int = 8, + num_replicas: int = 1, + rank: int = 0, +) -> tuple[DataLoader, DataLoader]: + """ + Factory function to create train and test dataloaders for RL prompt datasets. + + Args: + dataset_path: Path to dataset directory + dataset_type: "text" for TextPromptDataset or "geneval" for GenevalPromptDataset + split: Dataset split ("train" or "test") + train_batch_size: Batch size per GPU for training + test_batch_size: Batch size for testing + k: Number of times to repeat each sample (num_image_per_prompt) + seed: Random seed for sampler synchronization + train_num_workers: Number of workers for training dataloader + test_num_workers: Number of workers for test dataloader + num_replicas: Number of replicas (default 1 for single-node) + rank: Rank of current process (default 0 for single-node) + + Returns: + Tuple of (train_dataloader, test_dataloader) + """ + # Create datasets based on type + if dataset_type == "text": + train_dataset = TextPromptDataset(dataset_path, 'train') + test_dataset = TextPromptDataset(dataset_path, 'test') + collate_fn = TextPromptDataset.collate_fn + elif dataset_type == "geneval": + train_dataset = GenevalPromptDataset(dataset_path, 'train') + test_dataset = GenevalPromptDataset(dataset_path, 'test') + collate_fn = GenevalPromptDataset.collate_fn + else: + raise ValueError(f"Unknown dataset_type: {dataset_type}. Must be 'text' or 'geneval'") + + # Create infinite-loop training sampler + train_sampler = KRepeatSampler( + dataset=train_dataset, + batch_size=train_batch_size, + k=k, + num_replicas=num_replicas, + rank=rank, + seed=seed + ) + + # Create training dataloader with batch_sampler (infinite loop) + train_dataloader = DataLoader( + train_dataset, + batch_sampler=train_sampler, + num_workers=train_num_workers, + collate_fn=collate_fn, + ) + + # Create standard test dataloader + test_dataloader = DataLoader( + test_dataset, + batch_size=test_batch_size, + collate_fn=collate_fn, + shuffle=False, + num_workers=test_num_workers, + ) + + return train_dataloader, test_dataloader +