DataLoader sample by slices from Dataset

francois-rozet: <p>By default, <code>data.DataLoader</code> indexes elements of a batch one by one and collates them back into tensors. I have a dataset (subclass of <code>data.Dataset</code>) which can be indexed (efficiently) by slices. For example, the following can be indexed by slices:</p>
<pre><code class="lang-python">class MyDataset(data.Dataset):
   def __init__(self, a: Tensor, b: Tensor):
      self.a, self.b = a, b

   def __len__(self) -&gt; int:
        return len(self.a)

   def __getitem__(self, i) -&gt; Tuple[Tensor, Tensor]:
       return self.a[i], self.b[i]
</code></pre>
<p>I would like the <code>DataLoader</code> to take advantage of that. I think I need to use a <code>data.Sampler</code>, but I didn’t get at all how to from the documentation. Any idea ?</p>

ptrblck: <p>You could use a <a href="https://pytorch.org/docs/stable/data.html#torch.utils.data.BatchSampler"><code>BatchSampler</code></a>, which would pass multiple indices to the <code>Dataset.__getitem__</code> method.</p>

francois-rozet: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a>! To be sure, the following will send batches of shuffled indices to the dataset <code>__getitem__</code> ?</p>
<pre><code class="lang-python">dataset = MyDataset(a, b)
loader = data.DataLoader(
    dataset,
    batch_sampler=data.BatchSampler(data.RandomSampler(dataset), 256, False),
    num_workers=4,
    pin_memory=True
)
</code></pre>

ptrblck: <p>Yes, exactly!<br>
Here is a small code snippet to check the passed indices and the returned values:</p>
<pre><code class="lang-python">class MyDataset(Dataset):
    def __init__(self):
        self.data = torch.arange(100).float().view(-1, 1)
        
    def __getitem__(self, indices):
        print('dataset indices {}'.format(indices))
        x = self.data[indices]
        return x

    def __len__(self):
        return len(self.data)

dataset = MyDataset()
loader = DataLoader(
    dataset,
    # EDIT2: changed to sampler=...
    sampler=BatchSampler(RandomSampler(dataset), 2, False),
)

for x in loader:
    print('DataLoader loop {}'.format(x))
</code></pre>
<p>EDIT: wait, this looks wrong. Let me debug it quickly.</p>
<p>EDIT2: based on <a href="https://discuss.pytorch.org/t/how-to-use-batchsampler-with-getitem-dataset/78788/4">this code snippet</a> the indices would be directly passed to <code>__getitem__</code>, if <code>sampler=BatchSampler(...)</code> is used (which is also the case for my example code here).<br>
This seems to be an “edge case” maybe, as it would be similar to <a href="https://pytorch.org/docs/stable/data.html#disable-automatic-batching">disabling automatic batching</a>, but would use the sampler instead. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

francois-rozet: <p>Thanks again, <a class="mention" href="/u/ptrblck">@ptrblck</a>! I also had to use <code>batch_size=None</code> to prevent the collate function to add an additional dimension. So the final snippet is</p>
<pre><code class="lang-python">import torch
import torch.utils.data as data

class MyDataset(data.Dataset):
    def __init__(self):
        self.data = torch.rand(100, 5)

    def __len__(self) -&gt; int:
        return len(self.data)

    def __getitem__(self, i) -&gt; torch.Tensor:
        print(i)
        return self.data[i]

dataset = MyDataset()
sampler = data.BatchSampler(data.RandomSampler(dataset), 2, False)
loader  = data.DataLoader(dataset, sampler=sampler, batch_size=None)

for x in loader:
    print(x)
</code></pre>
<p>By the way, I feel like using the keyword <code>sampler</code> instead <code>batch_sampler</code> for specifically sampling batches is a bit weird. Perhaps, the interface could be improved a little ? Although, if it ain’t broke, don’t fix it <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=9" title=":sweat_smile:" class="emoji" alt=":sweat_smile:"></p>

ptrblck: <aside class="quote no-group" data-username="francois-rozet" data-post="5" data-topic="113005">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/francois-rozet/48/29950_2.png" class="avatar"> francois-rozet:</div>
<blockquote>
<p>By the way, I feel like using the keyword <code>sampler</code> instead <code>batch_sampler</code> for specifically sampling batches is a bit weird. Perhaps, the interface could be improved a little ? Although, if it ain’t broke, don’t fix it <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=12" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20"></p>
</blockquote>
</aside>
<p>Haha, I had the same feeling, but I’m sure I misunderstood it.</p>
<p>By checking the code again, the usage of <code>sampler=Sampler</code> and <code>batch_sampler=BatchSampler</code> would yield the same behavior inside the <code>Dataset</code>. The difference is only, that your <code>BatchSampler</code> can now yield multiple indices. This workflow can be used, if the <code>Dataset.__getitem__</code> method should stay the same, while you manipulate the sampler.</p>
<p>However, as shown in our code snippets, you could also pass a <code>BatchSampler</code> to the <code>sampler</code> argument, which would then pass <em>all</em> indices to the <code>__getitem__</code> method.<br>
While the naming is now a bit confusing, it’s a convenient method to load multiple samples in the <code>__getitem__</code>. I’m also not sure, if this is a “supported” use case, but so far it seems to work fine.</p>
<p><a class="mention" href="/u/vitalyfedyunin">@VitalyFedyunin</a> are we misusing the sampler or does it make sense?</p>