Question

一个人可以使用环境变量CUDA_VISIBLE_DEVICES来指定TensorFlow程序启动之前TensorFlow可以使用哪个GPU，例如CUDA_VISIBLE_DEVICES=1 python my_script.py。

我如何更改TensorFlow程序运行时可以使用的GPU？

我知道人们可以将程序更改为具有检查点，并使用CUDA_VISIBLE_DEVICES指定要使用的GPU从这些检查点重新运行TensorFlow程序，但是我想知道是否有某种方法可以避免依赖具有检查点的程序

Answer 1

考虑使用急切执行。它是针对您的网络进行这种动态调整的完美选择。

import random
import tensorflow as tf
tf.enable_eager_execution()

x = tf.constant(1)

for _ in range(10):
  dev = random.choice(['/cpu:0', '/gpu:0'])
  with tf.device(dev):
    x = x + 1
    print('res {} computed on {}'.format(x.numpy(), x.device))

在TensorFlow程序运行时，如何更改TensorFlow程序可以使用的GPU？

1 个答案: