Reputation: 6613
There are two ways I know of for controlling which GPUs a tf.Session
will take over. One is the environment variable CUDA_VISIBLE_DEVICES
, the other is initializing a tf.Session
with an explicit tf.ConfigProto
and setting config.gpu_options.visible_device_list
.
When using nccl
, the environment variable route doesn't work, so tf.ConfigProto
is the only way. However, my understanding is that in Tensorflow 2.0 tf.Session
construction is not explicit. Is there an equivalent of the ConfigProto.gpu_options.visible_device_list
route?
Upvotes: 3
Views: 9996
Reputation: 5689
In TF 2.0, the tf.ConfigProto
options have been moved in various namespace. Mainly tf.config and some in tf.debugging.
I also suggest to read the gpu guide.
For your question, I think this option was moved to: tf.config.experimental.set_visible_devices
Upvotes: 4
Reputation: 393
If you use the TF 2.0 upgrade script on your original TensorFlow model, tf.Session(config=...)
should change to tf.compat.v1.Session(config=...)
and work as expected. tf.compat.v1.ConfigProto
can be used to create the config parameter.
Upvotes: 4