Accumulating Batches for Contrastive Loss

myscience: <p>I have a custom dataset in which each example is fairly large (<code>batch, 80, 105, 90)</code>). I am training a self-supervised model with a contrastive loss that requires a decently large batch size (say <code>128</code>). My problem is that only <code>2</code> examples fit into GPU memory at once. However, before computing the loss, my model reduces the example to something of shape <code>(batch, 700)</code>. Does it make sense to <code>accumulate</code> these <em>latent</em> examples (which should fit into memory) and then compute my loss with a bigger <code>batch size</code>?  The idea would go something like this:</p>
<pre><code class="lang-python"># Training loop
bundle = (next(loader) for _ in range(accumulate))
latent = []
for pre_batch in bundle:
  latent += [model(pre_batch)]

latent = torch.cat(latent)
loss = contrastive_loss(latent)

optimizer.zero_grad()
loss.backward()
optimizer.step()
</code></pre>
<p>Based on this <a href="https://discuss.pytorch.org/t/why-do-we-need-to-set-the-gradients-manually-to-zero-in-pytorch/4903/20">post</a> my understanding is that this scenario is the “unlucky <span class="hashtag">#3</span>”, where both computation and memory are high (in fact, when I run this code I face a CUDA OOM error). Is there a smarter way of doing this or I a missing something trivial here?</p>

ptrblck: <aside class="quote no-group" data-username="myscience" data-post="1" data-topic="163453">
<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/myscience/48/53987_2.png" class="avatar"> myscience:</div>
<blockquote>
<p>Does it make sense to <code>accumulate</code> these <em>latent</em> examples (which should fit into memory) and then compute my loss with a bigger <code>batch size</code>?</p>
</blockquote>
</aside>
<p>It wouldn’t matter how small the output tensors are since the entire computation graph is stored in each iteration including all intermediates.</p>
<p>Yes, you are right that this approach will take a lot of memory, so would calling <code>backward</code> in each batch and accumulating the gradients work?</p>

myscience: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>! Thanks so much for the reply!</p>
<blockquote>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a><br>
Yes, you are right that this approach will take a lot of memory, so would calling <code>backward</code> in each batch and accumulating the gradients work?</p>
</blockquote>
<p>As I see it there is a problem in calling <code>backward</code> (I am assuming from <code>loss.backward()</code>) in each batch and simply start accumulating the gradients, but please correct me if this is nonsense. In this way I would compute each gradient with a <code>batch_size</code> of <code>2</code> given my GPU memory constraints, and because the contrastive loss (which is heavily inspired from the <code>CLIP</code> loss) depends on the <code>batch_size</code> not in a trivial averaging fashion (and the <em>quality</em> of the learning signal degrades for smaller sizes), I fear that in the standard gradient accumulation technique I would be averaging “very poor” gradients resulting in overall poor performances. This was the resoning that prompted the need for accumulating the batches and not simply the gradients.</p>
<p>For the time being I came up with a simple workaround: I put several <code>torch.utils.checkpoint.checkpoint</code> throughout the model and managed to fit a <code>batch_size</code> of <code>100</code>, which is close to what I was aiming for, of couse paying the additional computation cost.</p>

ptrblck: <p>I think your explanation makes sense and it’s good to hear checkpointing could help.</p>