U Ph,3@sddlmZddlmZmZmZmZddlmZddl m Z ddl m Z ddl mZddlmZddlmZdd lmZed \ZZGd d d e Z Gd ddee ZGddde ZdS)) annotations)CallableIterableIteratorSequence)Any)IterableDataset)get_worker_info)convert_tables_to_dicts)apply_transform) Randomizable)optional_importpandasc@s,eZdZdZd ddddddZd d ZdS) ra A generic dataset for iterable data source and an optional callable data transform when fetching a data sample. Inherit from PyTorch IterableDataset: https://pytorch.org/docs/stable/data.html?highlight=iterabledataset#torch.utils.data.IterableDataset. For example, typical input data can be web data stream which can support multi-process access. To accelerate the loading process, it can support multi-processing based on PyTorch DataLoader workers, every process executes transforms on part of every loaded data. Note that the order of output data may not match data source in multi-processing mode. And each worker process will have a different copy of the dataset object, need to guarantee process-safe from data source or DataLoader. Nz Iterable[Any]Callable | NoneNone)data transformreturncCs||_||_d|_dS)z Args: data: input data source to load and transform to generate dataset for model. transform: a callable data transform on input data. N)rrsource)selfrrrP/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/data/iterable_dataset.py__init__+szIterableDataset.__init__ccsvt}|dk r|jnd}|dk r&|jnd}t|j|_t|jD]0\}}|||kr@|jdk rjt|j|}|Vq@dS)Nr) r num_workersiditerrr enumeraterr )rinforriitemrrr__iter__5s    zIterableDataset.__iter__)N)__name__ __module__ __qualname____doc__rr!rrrrrs rcsZeZdZdZddddddfd d Zd d Zd dZfddZdddddZZ S) ShuffleBuffera Extend the IterableDataset with a buffer and randomly pop items. Args: data: input data source to load and transform to generate dataset for model. transform: a callable data transform on input data. buffer_size: size of the buffer to store items and randomly pop, default to 512. seed: random seed to initialize the random state of all workers, set `seed += 1` in every iter() call, refer to the PyTorch idea: https://github.com/pytorch/pytorch/blob/v1.10.0/torch/utils/data/distributed.py#L98. epochs: number of epochs to iterate over the dataset, default to 1, -1 means infinite epochs. Note: Both ``monai.data.DataLoader`` and ``torch.utils.data.DataLoader`` do not seed this class (as a subclass of ``IterableDataset``) at run time. ``persistent_workers=True`` flag (and pytorch>1.8) is therefore required for multiple epochs of loading when ``num_workers>0``. For example:: import monai def run(): dss = monai.data.ShuffleBuffer([1, 2, 3, 4], buffer_size=30, seed=42) dataloader = monai.data.DataLoader( dss, batch_size=1, num_workers=2, persistent_workers=True) for epoch in range(3): for item in dataloader: print(f"epoch: {epoch} item: {item}.") if __name__ == '__main__': run() Nrrintr) buffer_sizeseedepochsrcs,tj||d||_||_||_d|_dS)Nrrr)superrsizer*r+_idx)rrrr)r*r+ __class__rrrds zShuffleBuffer.__init__cCs4|t|||j|d}||j<||S)zAReturn the item at a randomized location `self._idx` in `buffer`.) randomizelenr/pop)rbufferretrrrrandomized_popkszShuffleBuffer.randomized_popccsNg}t|jD](}t||jkr,||V||q|rJ||Vq8dS)zLFill a `buffer` list up to `self.size`, then generate randomly popped items.N)rrr4r.r8append)rr6r rrr generate_itemrs  zShuffleBuffer.generate_itemc#s^|jd7_tj|jd|jdkr2t|jnttdD]}t||j dEdHq s         &H