Add option to run TensorFlow job on the preferred device

[zaleslaw]

System information

• TensorFlow version (you are using): 1.15 or 2.x
• Are you willing to contribute it (Yes/No): Yes

Describe the feature and the current behavior/state.
Now, we could run our code on GPU only via adding GPU dependencies to the classpath.
But the basic Python API provides an ability to set up the preferred device (GPU or CPU via device name)

The basic option also is available for low-level builder here

Will this change the current api? How?

Let's add the function tf.withDevice(“/GPU:0”) to the Scope class.

Who will benefit with this feature?
Anyone who trains neural network in distributed mode on different GPU/CPU devices.

Any Other info.

[karllessard]

Sounds good @zaleslaw !

I suggest that we go with a better typed API, maybe using some sort of enum, instead of hardcoded string like this (e.g. something like tf.withDevice(Device.GPU) or tf.withDevice(Device.GPU, 1), etc. Anyhow, something we can define together as we go, wdyt?

[zaleslaw]

I'm not very familiar with all possible device id strings, but it sounds good for me. пн, 22 июн. 2020 г. в 17:32, Karl Lessard <notifications@github.com>:

…

Sounds good @zaleslaw <https://github.com/zaleslaw> ! I suggest that we go with a better typed API, maybe using some sort of enum, instead of hardcoded string like this (e.g. something like tf.withDevice(Device.GPU) or tf.withDevice(Device.GPU, 1), etc. Anyhow, something we can define together as we go, wdyt? — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#77 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJEUHPT3LKD6ERYZRFK4MLRX5TO5ANCNFSM4OET2PHA> .

[Craigacp]

I'm in favour of an enum, and then an optional device ID number. We'll need it to specify what GPU something runs on, and I use shared GPU resources so I need to not contend with other users on different GPUs.

[zaleslaw]

Agree with Adam here: GPU/CPU mode with optional String IDs пн, 22 июн. 2020 г. в 18:43, Adam Pocock <notifications@github.com>:

…

I'm in favour of an enum, and then an optional device ID number. We'll need it to specify what GPU something runs on, and I use shared GPU resources so I need to not contend with other users on different GPUs. — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#77 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJEUHLZQEHI7JHWBCR4WSDRX533JANCNFSM4OET2PHA> .

[karllessard]

Yes, that was my proposal as well. I just want to make sure that we are not missing anything, I don't know if TensorFlow supports more options than simply /<device>:<id> in their string format, can you please double check @zaleslaw ?

[saudet]

The string also accepts host names and a lot of other things we'll have problems fitting in just enums:
https://www.tensorflow.org/guide/gpu
https://www.tensorflow.org/guide/tpu

[karllessard]

So probably some kind of builder could be more appropriate then? e.g. Device.job("localhost").replica(0).task(0).gpu(1)

[hthu]

Drive by comment.

Maybe look into existing python class?

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/framework/device_spec.py

[hthu]

I'm sending #83 for getting the DeviceSpec ported as a start.

[aday00]

Great idea! Happy to help test this with a multi-GPU system (#131).

Additional question: is there a command-line option to java to explicitly select which GPU should be used? GPU 0 is selected by default, but my system has 4 GPUs, so I'd like to run 4 java processes at once, with each java process explicitly selecting a different GPU.

[aday00]

Additional question: is there a command-line option to java to explicitly select which GPU should be used?

As expected, the CUDA_VISIBLE_DEVICES environment variable may be used to select GPU 1, for instance:

# export CUDA_VISIBLE_DEVICES=1; nice java -jar my.jar ...

[zaleslaw]

@aday00 and @karllessard looks like we have working device spec now, but have no ability to apply it for our graph or model (we have only setDevice option in GraphBuilder, but we could not pass it through the current TF operands and as a result we have no control on device yet.)

Please, correct me if I'm wrong.

[zaleslaw]

Maybe in eager mode it's better move to runner, but for the static graph look like we need to pre-define it earlier

[zaleslaw]

Yet one problem with imported graphs and models. When and how we could bind the device spec? Maybe after import we should have a method on graph, like addGradients? But in reality we need to rebuild our graph with new device specs.

I've ready to contribute in this ticket, but need some advices and tips before start

[karllessard]

Hey @zaleslaw , we should probably add the ability to add the DeviceSpec to the Scope when building an operation, similar to what we do with the operation name.

So we can end up doing something like:

tf.withName("myOp").withDevice(DeviceSpec.newBuilder()...).nn.avgPool(...)

If you feel that the syntax for building a device spec can still be improved, feel free to update it, thanks!