pytorch backend significantly slower than tensorflow

[ma-sadeghi]

I ran the "Poisson equation over L-shaped domain" example on an RTX 8000 once using tensorflow and once pytorch backends. Here's the timings:

tensorflow
----------
"adam" training ~ 90 seconds
"L-BFGS" training ~ 400 seconds

pytorch
-------
"adam" training ~ 1040 seconds
"L-BFGS" training ~ 530 seconds

which is ~ 10x slower for "adam" training. Is this expected? (and why) Thank you so much!

[lululxvi]

Tensorflow is accelerated by XLA by default. https://deepxde.readthedocs.io/en/latest/modules/deepxde.html#deepxde.config.disable_xla_jit

[ma-sadeghi]

Thank you for your reply. Is there a way to make deepxde faster with pytorch backend? (maybe not as fast as tensorflow, but at least in the same ballpark)

[lululxvi]

The current implementation is the same for all backends, and the algorithm is almost optimal. There might be some minor details that could be improved, but I don't think it will change anything significantly.

[lululxvi]

One possible big improvement is using PyTorch with XLA, but I don't have any experience on this.

[ma-sadeghi]

Thank you so much once again.