Is a GPU version of TF implied in order to run `live_demo.py` on a GPU?

[weirenorweiren]

I want to run live_demo.py on a GPU in Windows. The provided requirement files involves jax and tensorflow but live_demo.py only explicitly imports jax. Is a GPU version of tensorflow implied in order to run live_demo.py or jax on a GPU? Actually I am wondering why tensorflow is required in the requirements?

[cdoersch]

Tensorflow is required for some of the dataset loaders for training, but it shouldn't be required for inference (hence why it's not present in requirements_inference.txt or the base dependencies in or toml file). I don't think live_demo.py depends on tensorflow; let me know if I'm missing something.

Even for training, tensorflow shouldn't run on GPU; a CPU versoin should be fine.

[weirenorweiren]

Hi Carl, it seems live_demo.py requires tensorflow. Is there a way to circumvent this? If not, could you elaborate a bit more of why not tensorflow requiring a GPU version while jax is running on a GPU? I have no knowledge on this but intuitively I feel that if one thing (jax) runs on GPU and everything else would follow especially if it has a GPU version.

Traceback (most recent call last):
File "F:\Wei\tapnet\tapnet\live_demo.py", line 24, in
from tapnet.models import tapir_model
File "F:\Wei\tapnet\tapnet_init_.py", line 21, in
from tapnet.tapvid import evaluation_datasets # pylint:disable=g-importing-member
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "F:\Wei\tapnet\tapnet\tapvid\evaluation_datasets.py", line 33, in
import tensorflow as tf
ModuleNotFoundError: No module named 'tensorflow'

I care about this because of the below statement I found in https://www.tensorflow.org/install/source_windows. If GPU version for tensorflow is a must, I would need to decide whether using a lower version of TF or going for WSL2. Neither is preferred that ideally I hope I just need to install a wheel from https://github.com/cloudhan/jax-windows-builder.

Note: GPU support on native-Windows is only available for 2.10 or earlier versions, starting in TF 2.11, CUDA build is not supported for Windows. For using TensorFlow GPU on Windows, you will need to build/install TensorFlow in WSL2 or use tensorflow-cpu with TensorFlow-DirectML-Plugin

The discussion reminds me of one more thing. Are there any version restrictions for jaxlib and tensorflow? Here are the available wheels from the previous site: https://whls.blob.core.windows.net/unstable/index.html. It only supports a few versions of jaxlib. Would any of those versions be good for live_demo.py? And the same question applies for tensorflow.

[yangyi02]

Hi weirenorweiren,

Can you try comment this line

tapnet/tapnet/__init__.py

Line 21 in a0cd34d

from tapnet.tapvid import evaluation_datasets # pylint:disable=g-importing-member

then see if live_demo.py can work without tensorflow dependency?

Thanks.

[weirenorweiren]

Hi Yi, it works and thanks so much!

Could you also please take a look at the below question I asked in the previous reply? I am now working towards setting up the CUDA environment on Windows and any relevant information is appreciated!

The discussion reminds me of one more thing. Are there any version restrictions for jaxlib and tensorflow? Here are the available wheels from the previous site: https://whls.blob.core.windows.net/unstable/index.html. It only supports a few versions of jaxlib. Would any of those versions be good for live_demo.py? And the same question applies for tensorflow.

[yangyi02]

Hi @weirenorweiren , we are so far not very familiar with windows issues.

From what I can tell, if just for live_demo.py, you don't need tensorflow dependency.

For jaxlib, our code does not require any complex operations, so slightly early jaxlib version should also work.

We are also working on releasing a pytorch verion of live_demo.py, so that you will not need to bother about jaxlib and tensorflow.

[yangyi02]

Hi @weirenorweiren , we just released a pytorch live demo so jaxlib dependency is also not necessary any more. Feel free to give it a try.