How can I parallelize my Neural Network training

[oijoij]

Hello, new here. Sorry if I make any mistakes.

I am currently trying to train a Neural Network to approximate a function. However it is not fast enough, and I think the most obvious gains can come from parallelizing.

See this other discussion post to see the model I am trying to solve (though my code is different)

Solving HJB PDE using deep learning

Dear All, sorry for bothering with a possibly stupid question. I would like to ask for some advice about solving the Hamilton-Jacobi-Bellman equation using deep learning. I am trying to solve an HJB equation of a simple neoclassical growth model from macroeconomics. Its HJB equation is \rho V(k_t) = \frac{c_t^{1-\gamma}}{1-\gamma} + \frac{\partial V}{\partial k} \{A k_t^\alpha - \delta k_t - c_t\} Control variable in this problem is c_t \geq 0, characterized by following first-order conditi…

So a couple questions. First, how should I parallelize my training in order speed up my code the most? I plan on running K = {10,50,100, unsure} Neural Networks at a time. Should I parallelize by running a Neural Network on each core until it converges or reaches some max epoch ? Should I parallelize by distributing the “update step” (which I am doing via minibatch stochastic gradient descent and using ADAM). Should I parallelize by distributing the gradient step? Let me know if you have any thoughts.

Second of all, I am really struggling to figure out how to parallelize this. I have previously parallelized by using SharedArrays. However, SharedArrays needs to be of type bits and I am not really working with bits.

• So I am using ADAM optimizer in Flux. If my Neural Network updating step is done inside the parallelization, I need to make the optimizer a SharedArray I believe, but I can’t since it is of type ADAM.
• If I do my parallelization before I update my Neural Networks, then I’m thinking I’ll use it to grab all the gradients of my loss function. Well those are of type Zygote.Grads, so those can’t be in SharedArray form. If I convert those gradients to Array type, then I can store them in a SharedArray, but then I dont know how to use “Flux.Optimise.update!” since my gradient is no longer in the Zygote.Grads type.

I hope this makes sense! Please ask if it doesn’t and thank you for any help! Have spent quite some time on this and I am really struggling.