From af7dced53a6ef4ad891a8ff5bd2b41071d46cc23 Mon Sep 17 00:00:00 2001 From: Eitan Date: Mon, 23 Sep 2019 02:19:52 +0300 Subject: [PATCH] Expose num_workers in VideoClips --- torchvision/datasets/video_utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/torchvision/datasets/video_utils.py b/torchvision/datasets/video_utils.py index 4b81f64ec77..1d60b6c72db 100644 --- a/torchvision/datasets/video_utils.py +++ b/torchvision/datasets/video_utils.py @@ -49,8 +49,9 @@ class VideoClips(object): on the resampled video """ def __init__(self, video_paths, clip_length_in_frames=16, frames_between_clips=1, - frame_rate=None, _precomputed_metadata=None): + frame_rate=None, _precomputed_metadata=None, num_workers=1): self.video_paths = video_paths + self.num_workers = num_workers if _precomputed_metadata is None: self._compute_frame_pts() else: @@ -77,7 +78,7 @@ def __getitem__(self, idx): dl = torch.utils.data.DataLoader( DS(self.video_paths), batch_size=16, - num_workers=torch.get_num_threads(), + num_workers=self.num_workers, collate_fn=lambda x: x) with tqdm(total=len(dl)) as pbar: