Switch off the GPU in a running script in python

Viewed 786

The following code activates the GPU 0 and loads the necessary libraries like Cuda etc.

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

What about switching off the GPU in the running script when it is not needed any more? For example because the evaluation of a trained model needs to run on CPU. I tried to change the environment variable.

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

But this code does not switch off the GPU in the running script. It only can be used in the beginning of computing.

Any ideas are highly appreciated

Thanks in advance

1 Answers

There is not one single standard method for selecting GPUs in python, it will depend on the framework you are using in python to access the GPU.

The environment variable usage you describe is somewhat crude, and behaves as you describe. The environment variable affects CUDA behavior at the point of initialization of CUDA, which will typically only happen once, typically towards the beginning of your python script. Thereafter, changes to the environment variable will have no effect. CUDA samples it only once at initialization. In other respects, however, the effect of this should be consistent more-or-less independent of how you are using the GPU(s) in python.

To go beyond that, we need to talk about a specific framework. The underlying CUDA technology (CUDA C++) allows a framework developer to dynamically control which GPUs may be used at a particular point (from among those that are "exposed" via the environment variable; you cannot override that.) The method by which this dynamic control is exposed will vary based on framework. Since your question is tagged tensorflow, you can find a nice overview of how TF exposes GPU selection/control here.

Specifically you are asking about not using GPUs. I would refer you to the "Manual device placement" section. Note that "device" here refers to both CPU(s) and GPU(s), so the way to "shut off" GPU usage in this approach is to use the CPU device(s) instead, and examples are given there:

If you would like a particular operation to run on a device of your choice instead of what's automatically selected for you, you can use with tf.device to create a device context, and all the operations within that context will run on the same designated device.

tf.debugging.set_log_device_placement(True)

# Place tensors on the CPU
with tf.device('/CPU:0'):
  a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

# Run on the GPU
c = tf.matmul(a, b)
print(c)

Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

You will see that now a and b are assigned to CPU:0. Since a device was not explicitly specified for the MatMul operation, the TensorFlow runtime will choose one based on the operation and available devices (GPU:0 in this example) and automatically copy tensors between devices if required.

If you wanted the tf.matmul op to not use the GPU, it should only be necessary to move it into the scope of the preceding with statement.

Related