diff --git a/src/litdata/CHANGELOG.md b/src/litdata/CHANGELOG.md index d77da86a3..a08733fbc 100644 --- a/src/litdata/CHANGELOG.md +++ b/src/litdata/CHANGELOG.md @@ -10,6 +10,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). ### Fixed +- Pair each subsampled chunk with its corresponding region of interest in `train_test_split`. ([#949](https://github.com/Lightning-AI/litData/pull/949)) + - Fix `len(StreamingDataLoader)` undercounting per-worker partial batches with whole-chunk window shuffling. ([#939](https://github.com/Lightning-AI/litData/pull/939)) ## [0.2.76] - 2026-09-30 diff --git a/src/litdata/utilities/train_test_split.py b/src/litdata/utilities/train_test_split.py index 6e3413c27..14e4f2e08 100644 --- a/src/litdata/utilities/train_test_split.py +++ b/src/litdata/utilities/train_test_split.py @@ -50,8 +50,7 @@ def train_test_split( # we need subsampled chunk filenames, original chunk file, and subsampled_roi dummy_streaming_dataset = deepcopy_dataset(streaming_dataset) - # Note: We make this a set to speed up the computation of subsampled_chunks. - dummy_subsampled_chunk_filename = set(dummy_streaming_dataset.subsampled_files) + dummy_subsampled_chunk_filename = list(dummy_streaming_dataset.subsampled_files) dummy_subsampled_roi = dummy_streaming_dataset.region_of_interest subsampled_chunks: list[dict[str, Any]] = [] @@ -64,11 +63,16 @@ def train_test_split( # load chunks from `index.json` file data = load_index_file(input_dir.path) - # subsample the chunks defined in the index file + # subsample the chunks defined in the index file, keeping the dataset's own + # chunk order so that every chunk is paired with its own region of interest original_chunks = data["chunks"] - subsampled_chunks = [ - _org_chunk for _org_chunk in original_chunks if _org_chunk["filename"] in dummy_subsampled_chunk_filename - ] + chunk_by_filename = {_chunk["filename"]: _chunk for _chunk in original_chunks} + missing_chunk_filename = next( + (_filename for _filename in dummy_subsampled_chunk_filename if _filename not in chunk_by_filename), None + ) + if missing_chunk_filename is not None: + raise ValueError(f"Couldn't find chunk file {missing_chunk_filename} in the index file.") + subsampled_chunks = [chunk_by_filename[_filename] for _filename in dummy_subsampled_chunk_filename] new_datasets = [deepcopy_dataset(streaming_dataset) for _ in splits] diff --git a/tests/utilities/test_train_test_split.py b/tests/utilities/test_train_test_split.py index 9130ec902..85beb4afa 100644 --- a/tests/utilities/test_train_test_split.py +++ b/tests/utilities/test_train_test_split.py @@ -165,3 +165,35 @@ def test_streaming_dataset_subset(tmpdir): assert len(sliced) == 10 assert sliced[0] == 5 assert sliced[9] == 14 + + +def test_split_a_shuffled_subsampled_dataset(tmpdir): + cache = Cache(str(tmpdir), chunk_size=10) + for i in range(100): + cache[i] = i + cache.done() + cache.merge() + + ds = StreamingDataset(input_dir=str(tmpdir), subsample=0.45, shuffle=True, seed=1) + + train, test = train_test_split(ds, [0.6, 0.4], shuffle=False) + + # every split keeps the chunks and regions of interest of the dataset it splits + split_items = sorted([split[i] for split in (train, test) for i in range(len(split))]) + dataset_items = sorted(ds[i] for i in range(len(ds))) + assert split_items == dataset_items + + +def test_split_a_repeated_subsampled_dataset(tmpdir): + cache = Cache(str(tmpdir), chunk_size=10) + for i in range(100): + cache[i] = i + cache.done() + cache.merge() + + # a subsample above 1 repeats chunks, so the same filename is sampled more than once + ds = StreamingDataset(input_dir=str(tmpdir), subsample=1.5) + + train, test = train_test_split(ds, [0.5, 0.5]) + + assert len(train) + len(test) == len(ds) == 150