albanD:

What would be the feature?


Feature will-be native mini-batching support for torch.autograd.backward. Probably they could not be made parallel because of memory saving goal.



 albanD:

Also I’m curious what is the benefit of doing this compared to fw/bw or the smaller batches and accumulating the gradients?


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).



 albanD:

Note that you don’t have to do retain_grad, you can pass inputs= to backward to get the .grad field only for these Tensors.


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:
                    #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