Onednn: [Question] calculate batchnorm backward using forward output instead of forward input?

[eric-haibin-lin]

I wonder if it's possible to calculate the gradient of BN w.r.t weights based on the outputs of BN forward, instead of inputs of BN forward?

This way BN forward can release the memory for inputs as soon as forward is done, and those memory can be reused.

http://cthorey.github.io./backpropagation/

[emfomenk]

Hi @eric-haibin-lin,

Alas Intel MKL-DNN doesn't support this option.

[eric-haibin-lin]

Yes, I'm aware that it's not supported. Is there any plan to support this though?

[emfomenk]

No, we don't have plans to do so.

Could you please share what framework uses this approach?
So far we didn't get such a request from any of those...

[pengzhao-intel]

@emfomenk from Apache/MXNET :)

FYI, @eric-haibin-lin is the major architect of MXNET and co-working with Intel to improve the CPU perforamnce of MXNET.

[eric-haibin-lin]

In the backward pass, the gradient w.r.t β is straightforward.

With y, γ and β, we can calculate the value of ĥ. Now we have the gradient w.r.t. γ.

To calculate the gradient w.r.t h, we first need to reconstruct h. We can cache the mean and variance calculated in the forward pass, use that to reconstruct h, and finally calculate the gradient. This way we cache much smaller tensors (mean and var) compared to h.

We believe this will save lots of memory.

[eric-haibin-lin]

Actually on second thought, we cannot do inplace for both BN and Relu when these two are consecutive. We've enabled inplace optimization for Relu, which should be good enough.