This package extends DeepXDE to support the trick of Zero Coordinate Shift (ZCS, arxiv preprint), which can reduce GPU memory and wall time for training physics-informed DeepONets by an order of magnitude without affecting the training results.
Supported backends:
- Pytorch
- TensorFlow
- PaddlePaddle
pip install git+https://github.com/stfc-sciml/DeepXDE-ZCS
Note that the above line will upgrade your DeepXDE
to 1.10.0 or higher, as required by our extension.
Also, make sure you have one of the backends installed, which are not included in our requirements.
If you are familiar with DeepXDE, using DeepXDE-ZCS is straightforward with the following two steps.
Replacing the classes listed in the following table:
FROM | TO |
---|---|
deepxde.icbc.DirichletBC |
deepxde_zcs.DirichletBC_ZCS |
deepxde.icbc.IC |
deepxde_zcs.IC_ZCS |
deepxde.data.PDEOperatorCartesianProd |
deepxde_zcs.PDEOperatorCartesianProdZCS |
deepxde.Model |
deepxde_zcs.ModelZCS |
NOTE: Not all boundary conditions have been extended, and we are working on this. If you need one of them, please submit an issue to prioritize it. Doing this on your own is also easy, mostly by adding the function dimension to the outputs; see deepxde_zcs/boundary.py.
Take the diffusion-reaction equation for example. The original DeepXDE format reads like
import deepxde as dde
def pde(x, u, v):
d = 0.01
k = 0.01
u_t = dde.grad.jacobian(u, x, j=1)
u_xx = dde.grad.hessian(u, x, j=0)
return u_t - d * u_xx + k * u ** 2 - v
The DeepXDE-ZCS format looks similar, using our LazyGradZCS
class for derivative calculation:
import deepxde_zcs as ddez
def pde(zcs_parameters, u, v):
d = 0.01
k = 0.01
grad_zcs = ddez.LazyGradZCS(u, zcs_parameters)
u_t = grad_zcs.compute((0, 1))
u_xx = grad_zcs.compute((2, 0))
return u_t - d * u_xx + k * u ** 2 - v
In the above code:
- the input
zcs_parameters
is adict
generated by our code, and users don't need to care about what's in it; - the tuples
(0, 1)
and(2, 0)
give the wanted differential orders w.r.t.$(x,t)$ , i.e.,(0, 1)
for$\frac{\partial}{\partial t}$ and(2, 0)
for$\frac{\partial^2}{\partial x^2}$ ; - the
LazyGradZCS
class is smart enough to reuse any computed lower orders; for example, after computingu_xx
,u_x
will also be computed and cached, which can be given upon reqeust or serve as the starting point for computingu_xy
.
The diffusion-reaction equation is provided as a complete example under examples/diff_rec.
To run this example (change pytorch
to the backend you want):
- Aligned (original)
DDE_BACKEND=pytorch python diff_rec_aligned_pideeponet.py
- Unaligned (original)
DDE_BACKEND=pytorch python diff_rec_unaligned_pideeponet.py
- ZCS (ours)
DDE_BACKEND=pytorch python diff_rec_aligned_pideeponet_zcs.py
The GPU memory and wall time we measured on a Nvidia V100 are reported below. Note that this example is a small-scale problem for quick demo; ZCS can save more (in ratio) memory and time for larger-scale problems (e.g., those with more functions, more points, and higher-order PDEs).
-
PyTorch backend
METHOD GPU / MB TIME / s Aligned (original) 6611 165 Unaligned (original) 6785 160 ZCS (ours) 1487 11 -
TensorFlow backend
Our ZCS implementation is currently not jit-compiled. We are working on it for (maybe) faster running.
METHOD GPU / MB TIME / s Aligned (original) 9153 123 (jit) Unaligned (original) 7767 75 (jit) ZCS (ours) 847 34 (no jit) -
PaddlePaddle backend
METHOD GPU / MB TIME / s Aligned (original) 6423 207 Unaligned (original) 7541 374 ZCS (ours) 1353 16
Enjoy saving!