Average of Zygote gradients

Specific Domains Machine Learning
1

Hi Everyone,
I would like to know how I can manually calculate for example the mean over an array of gradients of a neural network. I have the following setup :

  • In the training loop I calculate the gradient of the loss
grad = Flux.gradient((m) -> abs2(y[i] - m(x[i])), net)[1]
  • at the end I want to have for example N gradients and compute the mean over these before i do the update of the network parameters. How can I perform this computation given that the gradients are computed as above.

Thanks.

2

Ok so I solved it with a naive approach by looping over the contents of the array and computing the mean, but is there any standard “Flux” way of computing this?

3

It sounds like you want Optimisers.AccumGrad

1 Like