Handling sparse batches to speed up training

username: czuczo_morales
cooked: <p>Hi everyone,</p>
<p>I’m working on a use case where the input to my model is a tensor of size <code>(B, K, context_length)</code>. Here:</p>
<ul>
<li><code>B</code> is the batch size,</li>
<li><code>K</code> represents different sources of information,</li>
<li><code>context_length</code> is the context length</li>
</ul>
<p>For any given sample, only a subset of these <code>K</code> sources contain meaningful information, while the rest are padded. When the sparsity in <code>K</code> is high (i.e., most sources are padded), this leads to inefficiencies in training due to unnecessary computations.</p>
<h3><a name="p-462736-proposed-idea-1" class="anchor" href="#p-462736-proposed-idea-1"></a>Proposed Idea</h3>
<p>To address this, I’m thinking of implementing a custom <code>DataLoader</code>. The idea is:</p>
<ol>
<li>For each sample, load only the non-padded (dense) parts of the <code>K</code> dimension.</li>
<li>While filling a batch, ensure that all sources for a single sample remain together (i.e., no scattering across batches). Continue sampling until the batch is complete.</li>
</ol>
<p>The goal is to avoid padded computations.</p>
<h3><a name="p-462736-questions-2" class="anchor" href="#p-462736-questions-2"></a>Questions</h3>
<ul>
<li>Has anyone implemented something similar, or is there a built-in PyTorch feature that could help with this?</li>
<li>Are there any suggestions or best practices for efficiently implementing such a custom <code>DataLoader</code>?</li>
</ul>

username: Tlotlo_Oepeng
cooked: <aside class="quote no-group" data-username="czuczo_morales" data-post="1" data-topic="215034">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/c/3ab097/48.png" class="avatar"> czuczo_morales:</div>
<blockquote>
<ul>
<li>Has anyone implemented something similar, or is there a built-in PyTorch feature that could help with this?</li>
<li>Are there any suggestions or best practices for efficiently implementing such a custom <code>DataLoader</code>?</li>
</ul>
</blockquote>
</aside>
<ul>
<li>To implement a custom dataloader, you would have to do it in cpp and then maybe bindings.</li>
<li>Consider using some kind of clever indexing. Or find someone who can do it for you. Knowledge of data structures is essential. if you can keep your lookups O(1), that would be amazing.</li>
</ul>

username: czuczo_morales
cooked: <p>Thank you for the contribution. The problem can be addressed easily using a <strong>custom Sampler</strong> that shuffles the dataset and composes batches <strong>before the epoch starts</strong>. The batching logic ensures that:</p>
<ol>
<li><strong>Samples remain coherent</strong>: All sources for a single sample are kept together within the same batch, avoiding scattering across batches.</li>
<li><strong>Batch size constraints are respected</strong>: The sum of non-zero entries in the <code>K</code> dimension across all samples in a batch does not exceed the specified <code>max_batch_size</code>.</li>
</ol>
<p>Zero entries can then be removed.</p>