pytorch_pfn_extras.distributed.create_distributed_subset_indices#

pytorch_pfn_extras.distributed.create_distributed_subset_indices(num_total_samples, num_replicas=None, rank=None, shuffle=True, seed=None)#

Returns a indices of a dataset to be used for the current process.

Parameters:

  • num_total_samples (int) – The size of the dataset.

  • num_replicas (Optional[int]) – Number of processes participating in the training. By default, torch.distributed.get_world_size() is used.

  • rank (Optional[int]) – Rank of the current process within num_replicas. By default, torch.distributed.get_rank() is used.

  • shuffle (bool) – If True (default), shuffle the indices.

  • seed (Optional[int]) – Random seed used to shuffle.

Return type:

List[int]