r"""Scatter method is used for splitting indices into subset, each subset will be assigned to a rank. Indices are evenly splitted by default. If customized indices assignment method is needed, please rewrite this method.
(self, indices)
| 109 | raise NotImplementedError |
| 110 | |
| 111 | def scatter(self, indices) -> List: |
| 112 | r"""Scatter method is used for splitting indices into subset, each subset |
| 113 | will be assigned to a rank. Indices are evenly splitted by default. |
| 114 | If customized indices assignment method is needed, please rewrite this method. |
| 115 | """ |
| 116 | total_size = self.num_samples * self.world_size |
| 117 | |
| 118 | # add extra indices to make it evenly divisible |
| 119 | indices += indices[: (total_size - len(indices))] |
| 120 | assert len(indices) == total_size |
| 121 | |
| 122 | # subsample |
| 123 | indices = indices[self.rank : total_size : self.world_size] |
| 124 | assert len(indices) == self.num_samples |
| 125 | |
| 126 | return indices |
| 127 | |
| 128 | def batch(self) -> Iterator[List[Any]]: |
| 129 | r"""Batch method provides a batch indices generator.""" |
no outgoing calls