Mini-batching / gradient accumulation within the model

vadimkantorov: <p>Paper <a href="https://arxiv.org/abs/2004.04725" class="inline-onebox" rel="noopener nofollow ugc">[2004.04725] Instance-aware, Context-focused, and Memory-efficient Weakly Supervised Object Detection</a> introduced <code>Sequential backprop</code> - within the model mini-batching of some components to make feasible handling of large number of ROIs with per-ROI ResNet’s res5 block.</p>
<p>I’ve made the following sketch to wrap ResNet’s res5 layer (haven’t tested yet).</p>
<p>In theory, can <code>torch.vmap</code> replace the mini-batching loop in <code>backward</code>?</p>
<p>Thanks!</p>
<pre><code class="lang-python">import torch
import torch.nn as nn

class SequentialBackprop(nn.Module):
    def __init__(self, module, batch_size = 1):
        super().__init__()
        self.module = module
        self.batch_size = batch_size

    def forward(self, x):
        y = self.module(x.detach())

        return self.Function.apply(x, y, self.batch_size, self.module)

    class Function(torch.autograd.Function):
        @staticmethod
        def forward(ctx, x, y, batch_size, module):
            ctx.save_for_backward(x)
            ctx.batch_size = batch_size
            ctx.module = module
            return y

        @staticmethod
        def backward(ctx, grad_output):
            (x,) = ctx.saved_tensors
            grads = []
            for x_mini, g_mini in zip(x.split(ctx.batch_size), grad_output.split(ctx.batch_size)):
                with torch.enable_grad():
                    x_mini = x_mini.clone()
                    x_mini.retain_grad()
                    y_mini = ctx.module(x_mini)
                torch.autograd.backward(y_mini, g_mini)
                grads.append(x_mini.grad)
            return torch.cat(grads), None, None, None

if __name__ == '__main__':
    backbone = nn.Linear(3, 6)
    neck = nn.Linear(6, 12)
    head = nn.Linear(12, 1)

    model = nn.Sequential(backbone, SequentialBackprop(neck, batch_size = 16), head)

    print('before', neck.weight.grad)

    x = torch.rand(512, 3)
    model(x).sum().backward()
    print('after', neck.weight.grad)

</code></pre>

vadimkantorov: <p><a class="mention" href="/u/alband">@alband</a> had to do nasty clone / retain_grad hacks, otherwise it complained about no grad accumulator of leaf tensor</p>

vadimkantorov: <p>Or maybe should this be a feature request for <code>torch.autograd.backward</code>?</p>

albanD: <blockquote>
<p>Or maybe should this be a feature request for torch.autograd.backward?</p>
</blockquote>
<p>What would be the feature?</p>
<p>Note that you don’t have to do retain_grad, you can pass <code>inputs=</code> to backward to get the .grad field only for these Tensors.</p>
<p>Also I’m curious what is the benefit of doing this compared to fw/bw or the smaller batches and accumulating the gradients?</p>

vadimkantorov: <aside class="quote no-group" data-username="albanD" data-post="4" data-topic="136460">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>What would be the feature?</p>
</blockquote>
</aside>
<p>Feature will-be native mini-batching support for torch.autograd.backward. Probably they could not be made parallel because of memory saving goal.</p>
<aside class="quote no-group" data-username="albanD" data-post="4" data-topic="136460">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>Also I’m curious what is the benefit of doing this compared to fw/bw or the smaller batches and accumulating the gradients?</p>
</blockquote>
</aside>
<p>It’s that you apply this gradient accumulation only in the middle of the model. If you do traditional gradient accumulation, you’d be underloading the GPU for the prior part of the model (the backbone).</p>
<aside class="quote no-group" data-username="albanD" data-post="4" data-topic="136460">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>Note that you don’t have to do retain_grad, you can pass <code>inputs=</code> to backward to get the .grad field only for these Tensors.</p>
</blockquote>
</aside>
<p>Not sure if I follow, I need gradients both for x_mini and for neck.weight - that’s why I can’t use just torch.autograd.grad. I’ve tried and got an error:</p>
<pre><code class="lang-auto">                    #x_mini.retain_grad()
                    y_mini = ctx.module(x_mini)
                torch.autograd.backward(y_mini, g_mini, inputs = x_mini)

RuntimeError: One of the differentiated Tensors given as 'inputs' to backward is not a leaf Tensor
</code></pre>

albanD: <aside class="quote no-group" data-username="vadimkantorov" data-post="5" data-topic="136460">
<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/vadimkantorov/48/365_2.png" class="avatar"> vadimkantorov:</div>
<blockquote>
<p>I’ve tried and got an error:</p>
</blockquote>
</aside>
<p>That should not happen on the latest version. We added that <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<blockquote>
<p>Feature will-be native mini-batching support for torch.autograd.backward. Probably they could not be made parallel because of memory saving goal.</p>
</blockquote>
<p>The thing is that this is not really an autograd thing, this is only true under a specific assumption that you have a batch of independent samples. Which is not true in general when not considering cummulative loss or(as soon as you have a batchnorm.<br>
The workaround of using a for-loop looks simple enough here.</p>

vadimkantorov: <aside class="quote no-group" data-username="albanD" data-post="4" data-topic="136460">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>to get the .grad field only for these Tensors.</p>
</blockquote>
</aside>
<p>Will it still accumulate gradients into self.module.weight or only produce gradInput?</p>
<p>Also, is there a way to bypass clone() / input copy?</p>

albanD: <blockquote>
<p>Will it still accumulate gradients into self.module.weight or only produce gradInput?</p>
</blockquote>
<p>Only the inputs given. But you can give <code>model.parameters()</code> as well <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>
<blockquote>
<p>Also, is there a way to bypass clone() / input copy?</p>
</blockquote>
<p>I’m not 100% sure why you need to clone here, I think detaching and requiring gradient again would work.</p>

vadimkantorov: <aside class="quote no-group" data-username="albanD" data-post="8" data-topic="136460">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>I’m not 100% sure why you need to clone here, I think detaching and requiring gradient again would work.</p>
</blockquote>
</aside>
<p>Thanks! Seems to work. Final version: <a href="https://gist.github.com/vadimkantorov/67fe785ed0bf31727af29a3584b87be1" class="inline-onebox" rel="noopener nofollow ugc">Mini-batching within the model in PyTorch · GitHub</a></p>

albanD: <p>Looks good. Note that the retain_grad is not needed any more as you already have a leaf Tensor that requires grad.</p>