Commit 333af7aa authored by James Thewlis's avatar James Thewlis Committed by Francisco Massa
Browse files

Move VideoClips dummy dataset to top level for pickling (#1649)

parent 4ab94cb7
......@@ -43,6 +43,21 @@ def unfold(tensor, size, step, dilation=1):
return torch.as_strided(tensor, new_size, new_stride)
class _DummyDataset(object):
"""
Dummy dataset used for DataLoader in VideoClips.
Defined at top level so it can be pickled when forking.
"""
def __init__(self, x):
self.x = x
def __len__(self):
return len(self.x)
def __getitem__(self, idx):
return read_video_timestamps(self.x[idx])
class VideoClips(object):
"""
Given a list of video files, computes all consecutive subvideos of size
......@@ -95,19 +110,9 @@ class VideoClips(object):
# strategy: use a DataLoader to parallelize read_video_timestamps
# so need to create a dummy dataset first
class DS(object):
def __init__(self, x):
self.x = x
def __len__(self):
return len(self.x)
def __getitem__(self, idx):
return read_video_timestamps(self.x[idx])
import torch.utils.data
dl = torch.utils.data.DataLoader(
DS(self.video_paths),
_DummyDataset(self.video_paths),
batch_size=16,
num_workers=self.num_workers,
collate_fn=lambda x: x)
......
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