Force DataLoader to fetch batched index from custom batch sampler

Antonio_Ossa: <p>Hi eveyone,</p>
<p>I’m working with a custom <code>Dataset</code> and <code>BatchSampler</code>. Due to the nature of my data, I have to fetch batches of different sizes, that’s why I’m using a <code>CustomBatchSampler</code>. Because of this, <code>DataLoader</code>s try to fetch items from my <code>CustomDataset</code> one item at each time.</p>
<ul>
<li>As you can see <a href="https://github.com/pytorch/pytorch/blob/a2d4d9eca6bfda3283bf2534260c0a8ba201b9e8/torch/utils/data/dataloader.py#L285" rel="nofollow noopener">here</a>, if I provide a <code>batch_sampler</code> to a <code>DataLoader</code>, <code>self.auto_collation</code> becomes <code>True</code>.</li>
<li>Then, because <code>self.auto_collation</code> is <code>True</code>, items are fetched one by one, as you can see <a href="https://github.com/pytorch/pytorch/blob/a2d4d9eca6bfda3283bf2534260c0a8ba201b9e8/torch/utils/data/_utils/fetch.py#L42" rel="nofollow noopener">here</a>.</li>
</ul>
<p>As a consequence, my code is faster if I iterate over my <code>CustomBatchSampler</code> and fetch multiple items manually, instead of using a <code>DataLoader</code>. I built my <code>Dataset</code> to support fetching multiple indexes at once. Also, I lose the possibility of using multiple workers easily (using a <code>DataLoader</code> as is).</p>
<p>Is it possible to force a <code>DataLoader</code> to fetch multiple items at once, like <code>data = self.dataset[possibly_batched_index]</code>? (like when <code>self.auto_collation</code> is <code>False</code>, in the second link). Is there a different approach expected to consider when using a <code>CustomBatchSampler</code> inside a <code>DataLoader</code>?</p>
<p>Thanks!</p>

ptrblck: <p>I think the <code>BatchSampler</code> will make sure to pass all batch indices to your <code>Dataset</code>'s <code>__getitem__</code> method as seen in this example:</p>
<pre><code class="lang-python">class MyDataset(Dataset):
    def __init__(self):
        self.data = torch.arange(100).view(100, 1).float()
        
    def __getitem__(self, index):
        print(index)
        x = self.data[index]
        return x
    
    def __len__(self):
        return len(self.data)

dataset = MyDataset()        
sampler = torch.utils.data.sampler.BatchSampler(
    torch.utils.data.sampler.RandomSampler(dataset),
    batch_size=10,
    drop_last=False)

loader = DataLoader(
    dataset,
    sampler=sampler)

for data in loader:
    print(data)
</code></pre>
<p>If you run it, you’ll see that the <code>index</code> inside the <code>Dataset.__getitem__</code> will contain 10 indices, which can be used to slice the data directly.</p>
<p>Let me know, if I misunderstood the question.</p>

Antonio_Ossa: <p>Thanks for your example, it worked perfectly and I managed to adapt it to my use case. My only problem is that I still can’t infer that I should do that from the <a href="https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader" rel="nofollow noopener">docs</a>, but maybe I’m not reading it properly. According to the docs, I understand that I have to provide my <code>CustomBatchSampler</code> instance as <code>batch_sampler</code>, not <code>sampler</code>, because it yields batches of indices. Is my understanding wrong or should the docs be improved (in case an issue about it is needed)?</p>
<p>Curiously, using a for loop on my <code>CustomBatchSampler</code> in the training loop takes almost the same amount of time than using the <code>DataLoader</code> adapted from your example. Is that expected? My guess is that its possible because what I’m doing manually, is the same as PyTorch does internally (in a <code>DataLoader</code>). Also, my dataset is an embedding already present in GPU (in my model), and I transfer the indexes through the forward method, so theres no much gain in using more workers.</p>
<ul>
<li>Train loop over <code>CustomBatchSampler</code> takes ~40s</li>
<li>Train loop over <code>DataLoader(num_workers=0)</code> takes ~40s</li>
<li>Train loop over <code>DataLoader(num_workers=2)</code> takes ~40s</li>
<li>Train loop over <code>DataLoader(num_workers=4)</code> takes ~40s</li>
</ul>
<p>(Machine has <code>os.cpu_cores()</code> = 2)</p>
<p>Could it be that the bottleneck is in the single GPU I’m using? Here’s a simplified sample of the code I’m using. I accumulate both accuracy and loss in CUDA Tensors to avoid transfer to/from CPU:</p>
<pre><code class="lang-auto"># CustomBatchSampler version
for data in train_batch_sampler:
    data = train_dataset[data]
    data_0 = torch.tensor(data[0], device=device)
    data_1 = torch.tensor(data[1], device=device)
    data_2 = torch.tensor(data[2], device=device)

    # Common section
    target = torch.ones(..., device=device)
    optimizer.zero_grad()
    with torch.set_grad_enabled(True):
        output = model(data_0, data_1, data_2)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
    running_acc.add_((output &gt; 0).sum())
    running_loss.add_(loss.detach() * output.size(0))
</code></pre>
<pre><code class="lang-auto"># DataLoader
for data in train_dataloader:
    data_0 = data[0].to(device, non_blocking=True).squeeze(dim=0)
    data_1 = data[1].to(device, non_blocking=True).squeeze(dim=0)
    data_2 = data[2].to(device, non_blocking=True).squeeze(dim=0)

    # Common section
    target = torch.ones(..., device=device)
    optimizer.zero_grad()
    with torch.set_grad_enabled(True):
        output = model(data_0, data_1, data_2)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
    running_acc.add_((output &gt; 0).sum())
    running_loss.add_(loss.detach() * output.size(0))
</code></pre>
<p>Thanks in advance <a class="mention" href="/u/ptrblck">@ptrblck</a></p>

ptrblck: <aside class="quote no-group" data-username="Antonio_Ossa" data-post="3" data-topic="86656">
<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/antonio_ossa/48/25953_2.png" class="avatar"> Antonio_Ossa:</div>
<blockquote>
<p>According to the docs, I understand that I have to provide my <code>CustomBatchSampler</code> instance as <code>batch_sampler</code> , not <code>sampler</code> , because it yields batches of indices. Is my understanding wrong or should the docs be improved (in case an issue about it is needed)?</p>
</blockquote>
</aside>
<p>I think I might have misunderstood the use case and might have given a wrong example.</p>
<p>The <code>batch_sampler</code> argument in the <code>DataLoader</code> will accept a sampler, which returns a batch of indices. Internally it will use the list comprehension (which you’ve linked to in the first post) and pass each index separately to <code>__getitem__</code>. This would make sure that the behavior of your custom <code>Dataset</code> can stay the same using the “standard” sampler and a <code>BatchSampler</code>.</p>
<p>However, we could just pass the <code>BatchSampler</code> as the <code>sampler</code> argument to the <code>DataLoader</code> (which I’ve done), so that we will get the complete batch of indices in <code>__getitem__</code>.</p>
<p>Let me know, if this is your use case or if I misunderstood it.</p>
<aside class="quote no-group" data-username="Antonio_Ossa" data-post="3" data-topic="86656">
<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/antonio_ossa/48/25953_2.png" class="avatar"> Antonio_Ossa:</div>
<blockquote>
<p>Curiously, using a for loop on my <code>CustomBatchSampler</code> in the training loop takes almost the same amount of time than using the <code>DataLoader</code> adapted from your example. Is that expected?</p>
</blockquote>
</aside>
<p>Besides the batching, shuffling, <code>drop_last</code> feature etc., the <code>DataLoader</code> might speed up your data loading pipeline, if you could load the data and process it in the background while your GPU is busy.<br>
If your data is already in the system memory and you could just slice it to create batches, you most likely won’t see a speedup (and might even slow down the code due to some overhead in the <code>DataLoader</code>).</p>

Antonio_Ossa: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="86656">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>However, we could just pass the <code>BatchSampler</code> as the <code>sampler</code> argument to the <code>DataLoader</code> (which I’ve done), so that we will get the complete batch of indices in <code>__getitem__</code> .</p>
</blockquote>
</aside>
<p>Yes, that was exactly the case and your example worked perfectly. The only difference with my initial code was that I passed my <code>CustomBatchSampler</code> instance as <code>batch_sampler</code>, not <code>sampler</code>, but as I said (thanks to your example) is fixed now.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="86656">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>If your data is already in the system memory and you could just slice it to create batches, you most likely won’t see a speedup (and might even slow down the code due to some overhead in the <code>DataLoader</code> ).</p>
</blockquote>
</aside>
<p>Well, that explains the timing I got. For the record, the times of my runs were almost the same iterating over both objects (<code>DataLoader</code> or <code>CustomBatchSampler</code>). The difference was visible when I used <code>num_workers &gt; 0</code> and added <code>torch.cuda.synchronize</code> at the end of each loop, which confirms that the <code>DataLoader</code> workers fetch the data while the GPU is busy.</p>
<p>Thanks a lot <a class="mention" href="/u/ptrblck">@ptrblck</a>!</p>

sogartar: <p>With this approach the <code>batch_size</code> in <code>DataLoader</code> gets defaulted to 1. The <code>DataLoader</code> will add an extra dimension of  size 1 to the loaded data.<br>
I found you could remove this by adding <code>batch_size=None</code> to the <code>DataLoader</code>.</p>
<pre><code class="lang-auto">loader = DataLoader(
    dataset,
    sampler=sampler,
    batch_size=None)
</code></pre>
<p>Then the <code>DataLoader</code> behaves similarly to when it does the batching itself, while retrieving one item at a time from the dataset.</p>
<pre><code class="lang-auto">loader = DataLoader(
    dataset,
    batch_size=10)
</code></pre>

Cow_woC: <p><a class="mention" href="/u/antonio_ossa">@Antonio_Ossa</a> I filed <a href="https://github.com/pytorch/pytorch/issues/71872" class="inline-onebox" rel="noopener nofollow ugc">Clarify the behavior of DataLoader sampler and batch_sampler parameters · Issue #71872 · pytorch/pytorch · GitHub</a> for the documentation issue. Feel free to comment there as needed.</p>