Give an option to either provide dataset or dataset_size in distributed sampler #1479
Conversation
facebook-github-bot
added
the
CLA Signed
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/data/1479
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 Cancelled JobAs of commit 1116f96 with merge base aef2409 ( CANCELLED JOB - The following job was cancelled. Please retry:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| @@ -179,19 +182,66 @@ def __iter__(self): | |||
| ) | |||
|
|
|||
|
|
|||
| class StatefulDistributedSampler(torch.utils.data.distributed.DistributedSampler): | |||
| class StatefulDistributedSampler(Sampler[int]): | |||
There was a problem hiding this comment.
I think we should continue subclassing DistributedSampler for StatefulDistributedSampler - it is easy to udnerstand that by just the naming and we might trigger many type checking issues in downstream code which uses StatefulDistributedSampler and expects a variant of DistributedSampler.
Since DistributedSampler is a common utility in PyTorch, StatefulDistributedSampler should be expected to be an extension of it.
There was a problem hiding this comment.
I decided to fork it instead of subclassing because I do not want to upstream these changes in torch.utils.data.distributed.DistributedSampler as it might break other users' code.
Nevertheless, it is redundant to have Dataset as an arg when we just need the length of it.
Currently the StatefulDistributedSampler takes the dataset as an argument, but only uses the length/size of the dataset.
Adding an option to provide size of the dataset instead, for more flexibility.