o # i@sdddlmZddlmZddlZddlmZddlmZddgZ GdddeZGd ddeZ dS) ) annotations)SequenceN)Dataset)DistributedSamplerr DistributedWeightedRandomSamplercs,eZdZdZ    ddfd d ZZS)ra Enhance PyTorch DistributedSampler to support non-evenly divisible sampling. Args: dataset: Dataset used for sampling. even_divisible: if False, different ranks can have different data length. for example, input data: [1, 2, 3, 4, 5], rank 0: [1, 3, 5], rank 1: [2, 4]. num_replicas: number of processes participating in distributed training. by default, `world_size` is retrieved from the current distributed group. rank: rank of the current process within `num_replicas`. by default, `rank` is retrieved from the current distributed group. shuffle: if `True`, sampler will shuffle the indices, default to True. kwargs: additional arguments for `DistributedSampler` super class, can be `seed` and `drop_last`. More information about DistributedSampler, please check: https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler. TNdatasetreven_divisiblebool num_replicas int | Nonerankshufflec  sptjd||||d||s6t|}||jkrtd|j|}|j||jkr1|jd8_||_dSdS)N)rr r r zBthe dataset length is less than the number of participating ranks.)super__init__lenr ValueError total_sizer num_samples) selfrrr r r kwargsdata_len extra_size __class__rU/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/data/samplers.pyr+s    zDistributedSampler.__init__)TNNT) rrrr r r r r r r )__name__ __module__ __qualname____doc__r __classcell__rrrrrscsBeZdZdZ     ddfdd ZfddZddZZS)ra  Extend the `DistributedSampler` to support weighted sampling. Refer to `torch.utils.data.WeightedRandomSampler`, for more details please check: https://pytorch.org/docs/stable/data.html#torch.utils.data.WeightedRandomSampler. Args: dataset: Dataset used for sampling. weights: a sequence of weights, not necessary summing up to one, length should exactly match the full dataset. num_samples_per_rank: number of samples to draw for every rank, sample from the distributed subset of dataset. if None, default to the length of dataset split by DistributedSampler. generator: PyTorch Generator used in sampling. even_divisible: if False, different ranks can have different data length. for example, input data: [1, 2, 3, 4, 5], rank 0: [1, 3, 5], rank 1: [2, 4].' num_replicas: number of processes participating in distributed training. by default, `world_size` is retrieved from the current distributed group. rank: rank of the current process within `num_replicas`. by default, `rank` is retrieved from the current distributed group. kwargs: additional arguments for `DistributedSampler` super class, can be `seed` and `drop_last`. NTrrweightsSequence[float]num_samples_per_rankr generatortorch.Generator | Nonerr r r c  sL|ddtjd||||d|||_|dur|n|j|_||_dS)Nr T)rrr r r) setdefaultrrr"rr$r%) rrr"r$r%rr r rrrrrXs  z)DistributedWeightedRandomSampler.__init__c#s\tt}tjfdd|Dtjd}tj|jdjd}|D]}||Vq$dS)Ncsg|]}j|qSr)r").0irrr ksz=DistributedWeightedRandomSampler.__iter__..)dtypeT)r%) listr__iter__torch as_tensordouble multinomialr$r%)rindicesr" rand_tensorr)rr*rr.is z)DistributedWeightedRandomSampler.__iter__cCs|jS)N)r$r*rrr__len__rsz(DistributedWeightedRandomSampler.__len__)NNTNN)rrr"r#r$r r%r&rr r r r r )rrrr rr.r5r!rrrrr@s  ) __future__rcollections.abcrr/torch.utils.datarrZ_TorchDistributedSampler__all__rrrrrs   )