Name	Size	Mode	Actions
backward_compatibility.cpython-36.pyc	474	0644	editdlrm
dataloader.cpython-36.pyc	24145	0644	editdlrm
dataloader_experimental.cpython-36.pyc	4284	0644	editdlrm
dataset.cpython-36.pyc	17051	0644	editdlrm
distributed.cpython-36.pyc	4851	0644	editdlrm
graph.cpython-36.pyc	1266	0644	editdlrm
sampler.cpython-36.pyc	9012	0644	editdlrm
sharding.cpython-36.pyc	1011	0644	editdlrm
_decorator.cpython-36.pyc	6031	0644	editdlrm
_typing.cpython-36.pyc	9852	0644	editdlrm
__init__.cpython-36.pyc	1477	0644	editdlrm

3 �Eg9�@sbdd
lZddl
mZmZmZ
dd
lZddlmZmZ
dd
l j Zeddd�ZGdd �d ee�Z d
S) �N)�TypeVar�Optional�Iterator�
)�Sampler�Dataset�T_coT)
� covariantc @sdeZ
dZd
Zdeeeeeeeedd�dd�Ze e d �
d d�Zed �
dd �Zedd�dd�Z dS)�DistributedSamplera�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. Args: dataset: Dataset used for sampling. num_replicas (int, optional): Number of processes participating in distributed training. By default, :attr:`world_size` 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``. .. 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. 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) NTr
F)�dataset�num_replicas�rank�shuffle�seed� drop_last�returncCs�|dkr tj
�std
�
�
tj�}|dkr@tj
�s8td
�
�
tj�}||ksP|dkrdtdj||d��
�
|
|_||_||_ d|_ ||_|jr�t|j�
|jdkr�t jt|j�
|j|j�
|_nt jt|j�
|j�
|_|j|j|_||_||_dS)Nz,Requires distributed package to be availabler
z7Invalid rank {}, rank should be in the interval [0, {}]r)�distZis_available�RuntimeErrorZget_world_sizeZget_rank� ValueError�formatrrr �epochr�len�math�ceil�num_samples� total_sizerr)�selfrrr rrr�r�H/usr/local/lib64/python3.6/site-packages/torch/utils/data/distributed.py�__init__;s0



















zDistributedSampler.__init__)
rc
Cs�|jr:t
j�}
|
j|j|j�

t
jt|j�
|
d
�j �}nt tt|j�
�
�
}|js�|j t|�
}|t|�
k
r~||d|�7}q�||tj|t|�
�
d|�7}n|d|j �}t|�
|j ks�t�
||j|j |j�}t|�
|jks�t�
t|�
S)N)
� generator)r�torch� GeneratorZmanual_seedrrZrandpermrr�tolist�list�rangerrrr�AssertionErrorr rr�iter)r�
g�indicesZpadding_sizerrr�__iter__^s




$

zDistributedSampler.__iter__c


Cs|jS)
N)
r)
rrrr�__len__ys
zDistributedSampler.__len__)rrcCs |
|_d
S)a)
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. N)
r)rrrrr� set_epoch|s zDistributedSampler.set_epoch)NNTr
F)�__name__� __module__�__qualname__�__doc__rr�int�boolrrrr*r+r,rrrrr s-

$!r )r�typingrrrr!�rrZtorch.distributedZdistributedrrr rrrr�
s