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

Leaving repeat aug sampler indices as tensor thrashes worker shared process memory

Browse Source
This commit is contained in:
Ross Wightman 2022-01-06 22:33:09 -08:00
parent 4df51f3932
commit 02ae11e526
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
timm/data/distributed_sampler.py
View File

@ -108,11 +108,11 @@ class RepeatAugSampler(Sampler):
indices = torch.arange(start=0, end=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....] # produce repeats e.g. [0, 0, 0, 1, 1, 1, 2, 2, 2....]
indices = torch.repeat_interleave(indices, repeats=self.num_repeats, dim=0) indices = torch.repeat_interleave(indices, repeats=self.num_repeats, dim=0).tolist()
# add extra samples to make it evenly divisible # add extra samples to make it evenly divisible
padding_size = self.total_size - len(indices) padding_size = self.total_size - len(indices)
if padding_size > 0: if padding_size > 0:
indices = torch.cat([indices, indices[:padding_size]], dim=0) indices += indices[:padding_size]
assert len(indices) == self.total_size assert len(indices) == self.total_size
# subsample per rank # subsample per rank