|
| 1 | +import math |
| 2 | +import torch |
| 3 | +from torch.utils.data import Sampler |
| 4 | +import torch.distributed as dist |
| 5 | +import torchvision.datasets.video_utils |
| 6 | + |
| 7 | + |
| 8 | +class DistributedSampler(Sampler): |
| 9 | + """ |
| 10 | + Extension of DistributedSampler, as discussed in |
| 11 | + https://github.com/pytorch/pytorch/issues/23430 |
| 12 | + """ |
| 13 | + |
| 14 | + def __init__(self, dataset, num_replicas=None, rank=None, shuffle=False): |
| 15 | + if num_replicas is None: |
| 16 | + if not dist.is_available(): |
| 17 | + raise RuntimeError("Requires distributed package to be available") |
| 18 | + num_replicas = dist.get_world_size() |
| 19 | + if rank is None: |
| 20 | + if not dist.is_available(): |
| 21 | + raise RuntimeError("Requires distributed package to be available") |
| 22 | + rank = dist.get_rank() |
| 23 | + self.dataset = dataset |
| 24 | + self.num_replicas = num_replicas |
| 25 | + self.rank = rank |
| 26 | + self.epoch = 0 |
| 27 | + self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas)) |
| 28 | + self.total_size = self.num_samples * self.num_replicas |
| 29 | + self.shuffle = shuffle |
| 30 | + |
| 31 | + def __iter__(self): |
| 32 | + # deterministically shuffle based on epoch |
| 33 | + g = torch.Generator() |
| 34 | + g.manual_seed(self.epoch) |
| 35 | + if self.shuffle: |
| 36 | + indices = torch.randperm(len(self.dataset), generator=g).tolist() |
| 37 | + else: |
| 38 | + indices = list(range(len(self.dataset))) |
| 39 | + |
| 40 | + # add extra samples to make it evenly divisible |
| 41 | + indices += indices[:(self.total_size - len(indices))] |
| 42 | + assert len(indices) == self.total_size |
| 43 | + |
| 44 | + # subsample |
| 45 | + indices = indices[self.rank:self.total_size:self.num_replicas] |
| 46 | + assert len(indices) == self.num_samples |
| 47 | + |
| 48 | + if isinstance(self.dataset, Sampler): |
| 49 | + orig_indices = list(iter(self.dataset)) |
| 50 | + indices = [orig_indices[i] for i in indices] |
| 51 | + |
| 52 | + return iter(indices) |
| 53 | + |
| 54 | + def __len__(self): |
| 55 | + return self.num_samples |
| 56 | + |
| 57 | + def set_epoch(self, epoch): |
| 58 | + self.epoch = epoch |
| 59 | + |
| 60 | + |
| 61 | +class UniformClipSampler(torch.utils.data.Sampler): |
| 62 | + """ |
| 63 | + Samples at most `max_video_clips_per_video` clips for each video, equally spaced |
| 64 | + Arguments: |
| 65 | + video_clips (VideoClips): video clips to sample from |
| 66 | + max_clips_per_video (int): maximum number of clips to be sampled per video |
| 67 | + """ |
| 68 | + def __init__(self, video_clips, max_clips_per_video): |
| 69 | + if not isinstance(video_clips, torchvision.datasets.video_utils.VideoClips): |
| 70 | + raise TypeError("Expected video_clips to be an instance of VideoClips, " |
| 71 | + "got {}".format(type(video_clips))) |
| 72 | + self.video_clips = video_clips |
| 73 | + self.max_clips_per_video = max_clips_per_video |
| 74 | + |
| 75 | + def __iter__(self): |
| 76 | + idxs = [] |
| 77 | + s = 0 |
| 78 | + # select at most max_clips_per_video for each video, uniformly spaced |
| 79 | + for c in self.video_clips.clips: |
| 80 | + length = len(c) |
| 81 | + step = max(length // self.max_clips_per_video, 1) |
| 82 | + sampled = torch.arange(length)[::step] + s |
| 83 | + s += length |
| 84 | + idxs.append(sampled) |
| 85 | + idxs = torch.cat(idxs).tolist() |
| 86 | + return iter(idxs) |
| 87 | + |
| 88 | + def __len__(self): |
| 89 | + return sum(min(len(c), self.max_clips_per_video) for c in self.video_clips.clips) |
0 commit comments