U Ph@sdddlmZddlmZddlZddlmZddlmZddgZ GdddeZGd ddeZ dS) ) annotations)SequenceN)Dataset)DistributedSamplerr DistributedWeightedRandomSamplercs0eZdZdZd ddddddfdd 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. TNrbool int | None)dataseteven_divisible num_replicasrankshufflec  sltjf||||d||sht|}||jkr:td|j|}|j||jkrb|jd8_||_dS)N)r r r r zBthe dataset length is less than the number of participating ranks.)super__init__lenr ValueError total_sizer num_samples) selfr r r r r kwargsdata_len extra_size __class__H/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/data/samplers.pyr+s   zDistributedSampler.__init__)TNNT)__name__ __module__ __qualname____doc__r __classcell__rrrrrs c sHeZdZdZddddddddd fd d Zfd d ZddZZS)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`. NTrzSequence[float]rztorch.Generator | Noner)r weightsnum_samples_per_rank generatorr r r c  sL|ddtjf||||d|||_|dk r:|n|j|_||_dS)Nr T)r r r r ) setdefaultrrr"rr#r$) rr r"r#r$r r r rrrrrXs z)DistributedWeightedRandomSampler.__init__c#sZtt}tjfdd|Dtjd}tj|jdjd}|D]}||VqFdS)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)rrrr rr,r3r!rrrrr@s  ) __future__rcollections.abcrr-torch.utils.datarrZ_TorchDistributedSampler__all__rrrrr s    )