Geoffrey Irving
Geoffrey Irving

Reputation: 6613

What's the equivalent of initializing a tf.Session with tf.ConfigProto in TensorFlow 2?

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

Answers (2)

Conchylicultor
Conchylicultor

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

dynamicwebpaige
dynamicwebpaige

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.

(GitHub issue.)

Upvotes: 4

Related Questions