multiple gpus

[Geotomo]

Dear Dr. Huang I am reading your code SA_PINN_Poisson2D.ipynb (Solving 2D Poisson Equation with SA-PINN). I download the code and it runs successfully, thanks for sharing your code. Do you have any idea of how to modify the script to use multiple gpus (using tf.distribute.MirroredStrategy())?

Best Regards Weigung

[huangboming]

Sorry for late reply.

TLDR: In order to use multiple GPUs for training, the first step is to modify the code to support GPU acceleration. After that, you can try multiple GPU.

Explainations:

In the code, there are two training steps, which are:

1. Adam training, which uses tf.keras.optimizers.Adam
2. L-BFGS training, which uses scipy.optimize.minimize

Although tf.keras.optimizers.Adam supports GPU acceleration, the problem is that scipy.optimize.minimize doesn't support GPU acceleration.

Maybe you can try tfp.optimizer.lbfgs_minimize for L-BFGS training. But TBH, I'm not sure if this works since I've tried to use it but failed(if my memory serves).

---

In order to use tf.distribute.MirroredStrategy(), all you need to do is to wrap the code of Train Model parts in the strategy.scope().

e.g.

strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    # the code of `Train Model` parts