TF 2.16.1 Fails to work with GPUs

[sgkouzias]

Hello everyone.

A respective issue is already open.

Description: It seems practically impossible for someone owning a PC with CUDA-enabled GPU to perform deep learning experiments with TensorFlow version 2.16.1 and utilize his GPU locally without manually performing some extra steps not included (until today) in the official TensorFlow documentation of the standard installation procedure of TensorFlow for Linux users with GPUs at least as a temporal fix! (it is a path issue)

Actions taken: I submitted a respective pull request in good faith and for the shake of all users as TensorFlow is "An Open Source Machine Learning Framework for Everyone". The request is still pending review (for one month until today).

Question: As Keras 3.0 is the default version for TensorFlow 2.16.1 do we have any updates regarding the bug? I wonder what other solutions exist.

[abhaskumarsinha]

Unless the error is fixed from the TensorFlow side, it is not possible to fix it with Keras as it is dependent on TensorFlow 2.16.1 ai its backend.

What you can do best is to utilize Keras distribution utilities and rewrite the code in Keras. Now you can switch to JAX and run distribution on CUDA GPU (if that is supported by JAX!).

Best Regards,
Abhas

[sgkouzias]

I am totally in favor of this option until the bug is fixed. Tested by training a simple model with just Keras code for the MNIST dataset powered by JAX as Keras backend and it was spectacularly fast and with no problems reported regarding my GPU utilization.

[abhaskumarsinha]

Sounds Good!

[sgkouzias]

Thank you very much.