mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Use torch.repeat_interleave() to generate repeated indices faster (#1058)

Browse Source 
* update: use numpy to generate repeated indices faster

* update: use torch.repeat_interleave() instead of np.repeat()

* refactor: remove unused import, numpy

* refactor: torch.range to torch.arange

* update: tensor to list before appending the extra samples

* update: concatenate the paddings with torch.cat
This commit is contained in:
Hyeongchan Kim 2022-01-03 07:01:06 +09:00 committed by GitHub
parent a79b2f3328
commit a0b2657497
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
timm/data/distributed_sampler.py
View File

@ -103,15 +103,16 @@ class RepeatAugSampler(Sampler):
g = torch.Generator()
g.manual_seed(self.epoch)
if self.shuffle:
indices = torch.randperm(len(self.dataset), generator=g).tolist()
indices = torch.randperm(len(self.dataset), generator=g)
else:
indices = list(range(len(self.dataset)))
indices = torch.arange(start=0, end=len(self.dataset))
# produce repeats e.g. [0, 0, 0, 1, 1, 1, 2, 2, 2....]
indices = [x for x in indices for _ in range(self.num_repeats)]
indices = torch.repeat_interleave(indices, repeats=self.num_repeats, dim=0)
# add extra samples to make it evenly divisible
padding_size = self.total_size - len(indices)
indices += indices[:padding_size]
if padding_size > 0:
indices = torch.cat([indices, indices[:padding_size]], dim=0)
assert len(indices) == self.total_size
# subsample per rank
@ -125,4 +126,4 @@ class RepeatAugSampler(Sampler):
return self.num_selected_samples
def set_epoch(self, epoch):
self.epoch = epoch
self.epoch = epoch