CODEWITHSUNDEEP
tensorflow
sudo-nim
sudo-nim

Reputation: 408

Get the device placement of a Tensorflow op programmatically

I want to write an assertion that ensures that certain ops in my graph will be run on a specific device. How can I determine programmatically the device placement of an op so that I can write such an assertion?

Upvotes: 2

Views: 1053

Answers (1)

BlueSun
BlueSun

Reputation: 3570

You can ensure that an operation is run on a specific device by using 

with tf.device('/gpu:0'):

before the definition of the operation (see here for more information).

edit:

Every available gpu has its own index: '/gpu:0', '/gpu:1', '/gpu:2' etc. This way you can bind specific operations to specific gpus. When you import tensorflow, it prints out which gpus are available and what index they are assigned.

(For example it prints: Creating TensorFlow device (/gpu:0) -> (device: 0, name: GeForce GTX 1070, pci bus id: ...))

Upvotes: 2

Related Questions