File size: 5,953 Bytes
8c212a5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 |
# python3.7
"""Contains the distributed data sampler.
This file is mostly borrowed from `torch/utils/data/distributed.py`.
However, sometimes, initialize the data loader and data sampler can be time
consuming (since it will load a large amount of data at one time). To avoid
re-initializing the data loader again and again, we modified the sampler to
support loading the data for only one time and then repeating the data loader.
Please use the class member `repeat` to control how many times you want the
data load to repeat. After `repeat` times, the data will be re-loaded.
NOTE: The number of repeat times should not be very large, especially when there
are too many samples in the dataset. We recommend to set `repeat = 500` for
datasets with ~50K samples.
"""
# pylint: disable=line-too-long
import math
from typing import TypeVar, Optional, Iterator
import torch
from torch.utils.data import Sampler, Dataset
import torch.distributed as dist
T_co = TypeVar('T_co', covariant=True)
class DistributedSampler(Sampler):
r"""Sampler that restricts data loading to a subset of the dataset.
It is especially useful in conjunction with
:class:`torch.nn.parallel.DistributedDataParallel`. In such a case, each
process can pass a :class:`~torch.utils.data.DistributedSampler` instance as a
:class:`~torch.utils.data.DataLoader` sampler, and load a subset of the
original dataset that is exclusive to it.
.. note::
Dataset is assumed to be of constant size.
Arguments:
dataset: Dataset used for sampling.
num_replicas (int, optional): Number of processes participating in
distributed training. By default, :attr:`rank` is retrieved from the
current distributed group.
rank (int, optional): Rank of the current process within :attr:`num_replicas`.
By default, :attr:`rank` is retrieved from the current distributed
group.
shuffle (bool, optional): If ``True`` (default), sampler will shuffle the
indices.
seed (int, optional): random seed used to shuffle the sampler if
:attr:`shuffle=True`. This number should be identical across all
processes in the distributed group. Default: ``0``.
drop_last (bool, optional): if ``True``, then the sampler will drop the
tail of the data to make it evenly divisible across the number of
replicas. If ``False``, the sampler will add extra indices to make
the data evenly divisible across the replicas. Default: ``False``.
current_iter (int, optional): Number of current iteration. Default: ``0``.
repeat (int, optional): Repeating number of the whole dataloader. Default: ``1000``.
.. warning::
In distributed mode, calling the :meth:`set_epoch` method at
the beginning of each epoch **before** creating the :class:`DataLoader` iterator
is necessary to make shuffling work properly across multiple epochs. Otherwise,
the same ordering will be always used.
"""
def __init__(self, dataset: Dataset, num_replicas: Optional[int] = None,
rank: Optional[int] = None, shuffle: bool = True,
seed: int = 0, drop_last: bool = False, current_iter: int = 0,
repeat: int = 1000) -> None:
super().__init__(None)
if num_replicas is None:
if not dist.is_available():
raise RuntimeError("Requires distributed package to be available")
num_replicas = dist.get_world_size()
if rank is None:
if not dist.is_available():
raise RuntimeError("Requires distributed package to be available")
rank = dist.get_rank()
self.dataset = dataset
self.num_replicas = num_replicas
self.rank = rank
self.iter = current_iter
self.drop_last = drop_last
# NOTE: self.dataset_length is `repeat X len(self.dataset)`
self.repeat = repeat
self.dataset_length = len(self.dataset) * self.repeat
if self.drop_last and self.dataset_length % self.num_replicas != 0:
# Split to nearest available length that is evenly divisible.
# This is to ensure each rank receives the same amount of data when
# using this Sampler.
self.num_samples = math.ceil(
(self.dataset_length - self.num_replicas) / self.num_replicas
)
else:
self.num_samples = math.ceil(self.dataset_length / self.num_replicas)
self.total_size = self.num_samples * self.num_replicas
self.shuffle = shuffle
self.seed = seed
self.__generate_indices__()
def __generate_indices__(self) -> None:
g = torch.Generator()
indices_bank = []
for iter_ in range(self.iter, self.iter + self.repeat):
g.manual_seed(self.seed + iter_)
indices = torch.randperm(len(self.dataset), generator=g).tolist()
indices_bank.extend(indices)
self.indices = indices_bank
def __iter__(self) -> Iterator[T_co]:
if self.shuffle:
# deterministically shuffle based on iter and seed
indices = self.indices
else:
indices = list(range(self.dataset_length))
if not self.drop_last:
# add extra samples to make it evenly divisible
indices += indices[:(self.total_size - len(indices))]
else:
# remove tail of data to make it evenly divisible.
indices = indices[:self.total_size]
# subsample
indices = indices[self.rank:self.total_size:self.num_replicas]
return iter(indices)
def __len__(self) -> int:
return self.num_samples
def __reset__(self, iteration: int) -> None:
self.iter = iteration
self.__generate_indices__()
# pylint: enable=line-too-long
|