torch.utils.data.distributed.DistributedSampler(dataset)

限制数据加载到数据集子集的采样器。
它与:class:`torch.nn.parallel.DistributedDataParallel`结合使用特别有用。
在这种情况下，每个进程都可以传递一个: class: ‘ ~ torch.utils.data的实例作为: class: ‘ ~ torch.utils.data.DataLoader`的Sampler，并加载原始数据集的独占子集。 

注意:
    Dataset是被指定固定大小。

参数Arg:
    dataset:Dataset被使用去采样.
    num_replicas(int, optional):参与分布式训练的进程数。默认情况下，从当前分布式组检索: attr:‘world _ size’。 
    rank(int, optional):Attr: ‘ num _ repicas’中当前进程的等级。默认情况下:attr: ‘Rank’从当前分布式组中检索。
    shuffle(bool, optional):如果是真(默认) ，采样器将洗牌的索引。
    seed(int, optional): 如果: attr: ‘ shuffle = True’,随机种子用于洗牌样本。在分布式组中的所有进程中，这个数字应该是相同的。默认值: “0”。
    drop_last (bool, optional): 如果为“ True”，则采样器将丢弃数据的尾部，使副本均匀可分。如果为“ False”，则采样器将添加额外的索引，以使数据在整个副本中均匀地可分

warning警告:
       分布式模式下，每一轮都调用: method: ‘ set _ epoch’方法，在每个epoch开始之前创建:class:`DataLoader`迭代器对于在多个epoch之间正确地进行洗牌是必要的。否则，将始终使用相同的顺序。

Example::

    >>> sampler = DistributedSampler(dataset) if is_distributed else None
    >>> loader = DataLoader(dataset, shuffle=(sampler is None),
    ...                     sampler=sampler)
    >>> for epoch in range(start_epoch, n_epochs):
    ...     if is_distributed:
    ...         sampler.set_epoch(epoch)
    ...     train(loader)
"""

源码

import math
from typing import TypeVar, Optional, Iterator

import torch
from . import Sampler, Dataset
import torch.distributed as dist


T_co = TypeVar('T_co', covariant=True)


class DistributedSampler(Sampler[T_co]):
    
    def __init__(self, dataset: Dataset, num_replicas: Optional[int] = None,
                 rank: Optional[int] = None, shuffle: bool = True,
                 seed: int = 0, drop_last: bool = False) -> 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()
        if rank >= num_replicas or rank < 0:
            raise ValueError(
                "Invalid rank {}, rank should be in the interval"
                " [0, {}]".format(rank, num_replicas - 1))
        self.dataset = dataset
        self.num_replicas = num_replicas
        self.rank = rank
        self.epoch = 0
        self.drop_last = drop_last
        # If the dataset length is evenly divisible by # of replicas, then there
        # is no need to drop any data, since the dataset will be split equally.
        if self.drop_last and len(self.dataset) % self.num_replicas != 0:  # type: ignore
            # 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(
                # `type:ignore` is required because Dataset cannot provide a default __len__
                # see NOTE in pytorch/torch/utils/data/sampler.py
                (len(self.dataset) - self.num_replicas) / self.num_replicas  # type: ignore
            )
        else:
            self.num_samples = math.ceil(len(self.dataset) / self.num_replicas)  # type: ignore
        self.total_size = self.num_samples * self.num_replicas
        self.shuffle = shuffle
        self.seed = seed

    def __iter__(self) -> Iterator[T_co]:
        if self.shuffle:
            # deterministically shuffle based on epoch and seed
            g = torch.Generator()
            g.manual_seed(self.seed + self.epoch)
            indices = torch.randperm(len(self.dataset), generator=g).tolist()  # type: ignore
        else:
            indices = list(range(len(self.dataset)))  # type: ignore

        if not self.drop_last:
            # add extra samples to make it evenly divisible
            padding_size = self.total_size - len(indices)
            if padding_size <= len(indices):
                indices += indices[:padding_size]
            else:
                indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size]
        else:
            # remove tail of data to make it evenly divisible.
            indices = indices[:self.total_size]
        assert len(indices) == self.total_size

        # subsample
        indices = indices[self.rank:self.total_size:self.num_replicas]
        assert len(indices) == self.num_samples

        return iter(indices)

    def __len__(self) -> int:
        return self.num_samples

    def set_epoch(self, epoch: int) -> None:
        r"""
        Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas
        use a different random ordering for each epoch. Otherwise, the next iteration of this
        sampler will yield the same ordering.

        Args:
            epoch (int): Epoch number.
        """
        self.epoch = epoch