From 41b14c6a1341cf84ec95878bcab5fa8e2b8a2fb7 Mon Sep 17 00:00:00 2001 From: sennnnn Date: Fri, 8 Mar 2024 22:54:17 +0800 Subject: [PATCH] Fix the memory error bug when loading metadata pickle file of vqvae dataset. Former-commit-id: 97ab243bd12915cea7aa2deecfac4b99417f78a3 --- opensora/models/ae/videobase/vqvae/dataset_vqvae.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/opensora/models/ae/videobase/vqvae/dataset_vqvae.py b/opensora/models/ae/videobase/vqvae/dataset_vqvae.py index c7f28c06a..13739319b 100644 --- a/opensora/models/ae/videobase/vqvae/dataset_vqvae.py +++ b/opensora/models/ae/videobase/vqvae/dataset_vqvae.py @@ -6,6 +6,7 @@ import warnings import glob import torch.utils.data as data +import torch.distributed as dist import torch.nn.functional as F from torchvision.datasets.video_utils import VideoClips @@ -41,7 +42,8 @@ class VQVAEDataset(data.Dataset): cache_file = osp.join(folder, f"metadata_{sequence_length}.pkl") if not osp.exists(cache_file): clips = VideoClips(files, sequence_length, num_workers=32) - pickle.dump(clips.metadata, open(cache_file, 'wb')) + if dist.is_initialized() and dist.get_rank() == 0: + pickle.dump(clips.metadata, open(cache_file, 'wb')) else: metadata = pickle.load(open(cache_file, 'rb')) clips = VideoClips(files, sequence_length,