Better compatibility with batch normalization

[andreped]

Better compatibility with batch normalization

[andreped]

I implemented a custom batch norm layer recently, which was made available in v0.3.2.

Implementation can be seen here.

The plan is to add accumulation to the call step similar as done in the Model wrapper.

The performance on the custom layer reaches very similar results to keras' BN layer, but we should get identical before we can say that it is working as intended.

[andreped]

To add accumulation support to the custom BN layer, relevant resources are this, this, and this.

[andreped]

I have added gradient accumulation support to the custom BN layer now in 05fb499.

However, unit tests show that the results are not equivalent when increasing accum_steps. Hence, it is not yet working as intended.

[andreped]

Support for gradient accumulation to batch normalization has been added. Even though it is not perfect, at least people could test it and improve it further.

I also updated the documentations regarding how to use it in f54e389.

We can open a separate issue to benchmark it properly. Might also be that opening a discussion in the Discussions tab is the way to go. Hence, closing this issue.