I highly doubt that it works well when removing this.

The idea of clearing the gradients at that point is that the training of each minibatch starts out with a clean slate. If we don't do that, the gradients accumulate over multiple minibatches leading to unexpected behaviour.

As an alternative the clearing of the weights could be moved to the end of the function, but that would make some diagnostics like looking at the gradients after each minibatch hard/impossible. Personally I prefer the current way between the two alternatives.

Hmm, I was under impression that weight gradients for i-th layer are recomputed in the backpropagation pass from the following data:

• output of i-th layer (computed during forward pass),
• output gradients of i-th layer (computed in the previous backpropagation step for i+1-th layer or provided by solver for the last layer),
• weights of i-th layer if layer has weigths (computed in the previous epoch),

So weight gradients for i-th are strictly output of the computation, so their contents before computation don't matter.

I've also tried to verify it in practice, here are graphs for mlp and conv with and without this patch. It looks like there is no visible change.

Edit: changed wording a bit.

Hey @alexandermorozov, thanks for the PR and the charts 👍

Max and I agree, that it makes sense to merge the PR as it provides immediate value. If we should decide later to handle the gradients differently, it would be quite easy to put the one line back in.

I am happy to merge!

@homu r+

📌 Commit 01be854 has been approved by MichaelHirn

⌛ Testing commit 01be854 with merge 434a0db...

💥 Test timed out

@MichaelHirn, thanks! It's great that this PR is useful )

I've rebased this branch, it should merge now.

@homu r+

📌 Commit 6c4482c has been approved by hobofan

⚡ Test exempted - status