squirrel.iterstream.torch_composables
¶
Module Contents¶
Classes¶
Composable to split data between ranks of a multi-rank loading setup |
|
Composable to split data between PyTorch workers of a single rank |
|
Mixin-Composable to have squirrel pipeline inherit from PyTorch IterableDataset |
Functions¶
|
Returns a callable that takes an iterable and applies a skipping operation on it. |
Attributes¶
-
squirrel.iterstream.torch_composables.
logger
¶
-
class
squirrel.iterstream.torch_composables.
SplitByRank
(source: Optional[Iterable] = None, torch_dist_group: Optional[str] = None)¶ Bases:
squirrel.iterstream.base.Composable
Composable to split data between ranks of a multi-rank loading setup
Init the SplitByRank composable.
-
__iter__
(self) → Iterator¶ Method to iterate over the source and yield the elements that will be processed by a particular node
-
-
class
squirrel.iterstream.torch_composables.
SplitByWorker
(source: Optional[Iterable] = None)¶ Bases:
squirrel.iterstream.base.Composable
Composable to split data between PyTorch workers of a single rank
Init
-
__iter__
(self) → Iterator¶ Method to iterate over the source and yield the elements that will be processed by a particular worker
-
-
class
squirrel.iterstream.torch_composables.
TorchIterable
(source: Optional[Iterable] = None)¶ Bases:
squirrel.iterstream.base.Composable
,torch.utils.data.IterableDataset
Mixin-Composable to have squirrel pipeline inherit from PyTorch IterableDataset
Init
-
__iter__
(self) → Iterator¶ Method to iterate over the source
-
-
squirrel.iterstream.torch_composables.
skip_k
(rank: int, world_size: int) → Callable[[Iterable], Iterator]¶ Returns a callable that takes an iterable and applies a skipping operation on it.
- Parameters
rank – int denoting the rank of the distributed training process.
world_size – int denoting the full world size.