(self, dataset, given_size=None, dup_shuffle=False, world_size=None, rank=None)
| 817 | |
| 818 | class DistributedGivenSizeSampler(Sampler): |
| 819 | def __init__(self, dataset, given_size=None, dup_shuffle=False, world_size=None, rank=None): |
| 820 | if world_size is None: |
| 821 | world_size = get_world_size() |
| 822 | if rank is None: |
| 823 | rank = get_rank() |
| 824 | assert rank < world_size |
| 825 | self.dataset = dataset |
| 826 | self.dup_shuffle = dup_shuffle |
| 827 | self.world_size = world_size |
| 828 | self.rank = rank |
| 829 | self.epoch = 0 |
| 830 | if given_size is None: |
| 831 | self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.world_size)) |
| 832 | else: |
| 833 | self.num_samples = int(math.ceil(given_size * 1.0 / self.world_size)) |
| 834 | self.total_size = self.num_samples * self.world_size |
| 835 | |
| 836 | if self.dup_shuffle: |
| 837 | self.offset = 0 |
| 838 | self.g = torch.Generator() |
| 839 | #self.g.manual_seed(self.rank) |
| 840 | self.indices = self.gen_new_list(self.g) |
| 841 | |
| 842 | def __iter__(self): |
| 843 | if self.dup_shuffle: |
nothing calls this directly
no test coverage detected