How can I combine tf_federated with tensorflow-gpu

General Discussion
, ,
#1

I installed tensorflow-gpu 2.4.1 by the command ‘conda install tensorflow-gpu’, and it could recognize my gpu, that is, it will return True when I call tf.test.is_gpu_available().
But when I was installing tensorflow_federated by ‘pip install tensorflow_federated’, I found that my tensorflow-gpu installed by conda was uninstalled. And after the pip installation, tensorflow could not recognize my gpu anymore.
How can I use tensorflow_federated on GPU?

#3

Hi Jason,

What I think might have happened is that you installed the version 0.19 of TF Federated and that depends on TF 2.5 and it might have updated your 2.4 installation. (versions table: GitHub - tensorflow/federated: A framework for implementing federated learning)

You can verify these by

print(tf.__version__) 
print(tff.__version__)

on your env.

If you install TF 2.5-gpu everything should work