From 0b3d2db9f319617540fe86d1148b3fba24c803c6 Mon Sep 17 00:00:00 2001 From: Raghav K Date: Thu, 5 Nov 2026 18:31:54 -0800 Subject: [PATCH] [bugfix] Stop the preprocessing dataset from keeping every decoded video in memory (#2069) Co-authored-by: Claude Opus 5.5 Co-authored-by: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> --- fastvideo/dataset/preprocessing_datasets.py | 4 +- .../test_preprocessing_dataset_memory.py | 41 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) create mode 100644 fastvideo/tests/dataset/test_preprocessing_dataset_memory.py diff --git a/fastvideo/dataset/preprocessing_datasets.py b/fastvideo/dataset/preprocessing_datasets.py index 888aae21d..9e27f86b4 100644 --- a/fastvideo/dataset/preprocessing_datasets.py +++ b/fastvideo/dataset/preprocessing_datasets.py @@ -5,7 +5,7 @@ import os import random from abc import ABC, abstractmethod from collections import Counter -from dataclasses import dataclass +from dataclasses import dataclass, replace from os.path import join as opj from typing import Any @@ -513,7 +513,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, torch.distribu def _get_item(self, idx: int) -> dict: """Get a single processed data item.""" - batch = self.processed_batches[idx] + batch = replace(self.processed_batches[idx]) # Apply transformation stages batch = self.video_transform_stage.process(batch) diff --git a/fastvideo/tests/dataset/test_preprocessing_dataset_memory.py b/fastvideo/tests/dataset/test_preprocessing_dataset_memory.py new file mode 100644 index 000000000..a3663b91e --- /dev/null +++ b/fastvideo/tests/dataset/test_preprocessing_dataset_memory.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VideoCaptionMergedDataset must not keep decoded pixels on its stored metadata.""" + +import argparse +import json + +import torch + +from fastvideo.dataset import getdataset + + +def _dataset(tmp_path, clips: int): + videos = tmp_path / "videos" + videos.mkdir() + annotations = [] + for i in range(clips): + (videos / f"clip{i}.mp4").write_bytes(b"") + annotations.append({"path": f"clip{i}.mp4", "cap": ["a caption"], "resolution": {"width": 32, "height": 32}, + "fps": 8, "duration": 1.0, "num_frames": 8}) + (tmp_path / "annotation.json").write_text(json.dumps(annotations)) + (tmp_path / "merge.txt").write_text(f"{videos},{tmp_path / 'annotation.json'}\n") + args = argparse.Namespace(model_path=str(tmp_path / "no-model"), cache_dir=str(tmp_path), num_frames=8, + train_fps=8, speed_factor=1.0, video_length_tolerance_range=5, drop_short_ratio=1.0, + text_max_length=32, training_cfg_rate=0.0, max_height=32, max_width=32, + do_temporal_sample=False, data_merge_path=str(tmp_path / "merge.txt"), seed=0) + return getdataset(args) + + +def test_iterating_does_not_retain_pixels_on_processed_batches(tmp_path, monkeypatch) -> None: + dataset = _dataset(tmp_path, clips=3) + + def fake_decode(batch): + batch.pixel_values = torch.zeros(3, 8, 32, 32) + return batch + + monkeypatch.setattr(dataset.video_transform_stage, "process", fake_decode) + + items = list(dataset) + + assert len(items) == 3 and all(item["pixel_values"] is not None for item in items) + assert all(batch.pixel_values is None for batch in dataset.processed_batches)