Skip to content

GitLab

Commit 02a8c0ad authored by ekosman's avatar ekosman Committed by Francisco Massa
Browse files

Expose num_workers in VideoClips (#1359)

parent 76c04d68
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
torchvision/datasets/video_utils.py
View file @ 02a8c0ad
...@@ -49,8 +49,9 @@ class VideoClips(object): ...@@ -49,8 +49,9 @@ class VideoClips(object):
on the resampled video on the resampled video
""" """
def __init__(self, video_paths, clip_length_in_frames=16, frames_between_clips=1, 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.video_paths = video_paths
self.num_workers = num_workers
if _precomputed_metadata is None: if _precomputed_metadata is None:
self._compute_frame_pts() self._compute_frame_pts()
else: else:
...@@ -77,7 +78,7 @@ class VideoClips(object): ...@@ -77,7 +78,7 @@ class VideoClips(object):
dl = torch.utils.data.DataLoader( dl = torch.utils.data.DataLoader(
DS(self.video_paths), DS(self.video_paths),
batch_size=16, batch_size=16,
num_workers=torch.get_num_threads(), num_workers=self.num_workers,
collate_fn=lambda x: x) collate_fn=lambda x: x)
with tqdm(total=len(dl)) as pbar: with tqdm(total=len(dl)) as pbar:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment