jaxopt: Levenberg-Mardquat Running exceptionally slow unless verbose is enabled

Levenberg-Mardquat optimizer runs exceptionally slow (~30 seconds for 30 iterations) until I turn on verbose==True (~1 second for 30 iterations. Any idea what may be going on? enabling JIT seems to have no impact. Was hoping to use this for a real-time system but even at 1 second things are way too slow.

About this issue

Original URL
State: closed
Created 2 years ago
Comments: 23

Most upvoted comments

I looked at your colab more closely and think that the timings need to be performed more carefully. In particular, you need to make sure that you are not averaging the compilation time with running time when using jaxopt’s LM. I rearranged some of your timing code and used time.process_time to extract the CPU time. What I see is that, the compilation of both approaches (jaxopt original implementation and the custom version) are similar. Jaxopt’s runtime is actually being returned much less if you make sure you are not including the compilation time. Let me know if you can reproduce this on your end.

amir-saadat on Aug 3, 2022