probability icon indicating copy to clipboard operation
probability copied to clipboard

Importing substrates is slow

Open carlosgmartin opened this issue 2 years ago • 2 comments

It takes a noticeably long time to just import TFP substrates:

$ time python3 -c "import tensorflow_probability.substrates.jax"
python3 -c "import tensorflow_probability.substrates.jax"  4.72s user 0.77s system 116% cpu 4.724 total

$ time python3 -c "import tensorflow_probability.substrates.numpy"
python3 -c "import tensorflow_probability.substrates.numpy"  4.75s user 0.86s system 115% cpu 4.843 total

Might it be possible to speed up these imports somehow?

carlosgmartin avatar Apr 17 '23 20:04 carlosgmartin

Do you have tensorflow installed? Uninstalling that might speed up the import, if you're only going to use JAX or numpy. It would be nice to disentangle how much of the import time comes from importing the respective lib (jax | numpy). I would think that would account for a decent chunk. We could push lazy module loading deeper into the library, if it's actually beneficial, but we'd want to see some more precise measurements.

Brian Patton | Software Engineer | @.***

On Mon, Apr 17, 2023 at 4:17 PM Carlos Martin @.***> wrote:

It takes a noticeably long time to just import TFP substrates:

$ time python3 -c "import tensorflow_probability.substrates.jax" python3 -c "import tensorflow_probability.substrates.jax" 4.72s user 0.77s system 116% cpu 4.724 total

$ time python3 -c "import tensorflow_probability.substrates.numpy" python3 -c "import tensorflow_probability.substrates.numpy" 4.75s user 0.86s system 115% cpu 4.843 total

Might it be possible to speed up these imports somehow?

— Reply to this email directly, view it on GitHub https://github.com/tensorflow/probability/issues/1714, or unsubscribe https://github.com/notifications/unsubscribe-auth/AFJFSIYSG2WPZ5ECJ5QXIETXBWQNZANCNFSM6AAAAAAXBVVAXI . You are receiving this because you are subscribed to this thread.Message ID: @.***>

brianwa84 avatar May 02 '23 15:05 brianwa84

@brianwa84 Uninstalling tensorflow indeed reduced the time. With tensorflow:

python3 -c "import tensorflow_probability.substrates.jax"  5.85s user 1.28s system 90% cpu 7.858 total
python3 -c "import tensorflow_probability.substrates.numpy"  5.16s user 0.94s system 109% cpu 5.590 total

Without tensorflow:

python3 -c "import tensorflow_probability.substrates.jax"  1.57s user 0.35s system 153% cpu 1.254 total
python3 -c "import tensorflow_probability.substrates.numpy"  1.12s user 0.27s system 241% cpu 0.575 total

carlosgmartin avatar May 02 '23 19:05 carlosgmartin