| 25 | |
| 26 | |
| 27 | class StatefulDistributedSampler(DistributedSampler): |
| 28 | |
| 29 | def __init__( |
| 30 | self, |
| 31 | dataset: Dataset, |
| 32 | num_replicas: Optional[int] = None, |
| 33 | rank: Optional[int] = None, |
| 34 | shuffle: bool = True, |
| 35 | seed: int = 0, |
| 36 | drop_last: bool = False, |
| 37 | ) -> None: |
| 38 | super().__init__(dataset, num_replicas, rank, shuffle, seed, drop_last) |
| 39 | self.start_index: int = 0 |
| 40 | |
| 41 | def __iter__(self) -> Iterator: |
| 42 | iterator = super().__iter__() |
| 43 | indices = list(iterator) |
| 44 | indices = indices[self.start_index:] |
| 45 | return iter(indices) |
| 46 | |
| 47 | def __len__(self) -> int: |
| 48 | return self.num_samples - self.start_index |
| 49 | |
| 50 | def reset(self) -> None: |
| 51 | self.start_index = 0 |
| 52 | |
| 53 | def state_dict(self, step) -> dict: |
| 54 | return {'start_index': step} |
| 55 | |
| 56 | def load_state_dict(self, state_dict: dict) -> None: |
| 57 | self.__dict__.update(state_dict) |
| 58 | |
| 59 | |
| 60 | class VariableVideoBatchSampler(DistributedSampler): |