tensorflow - 如何在tensorflow中设置特定的gpu？

Question

我想指定 gpu 来运行我的进程。我将其设置如下:

import tensorflow as tf
with tf.device('/gpu:0'):
    a = tf.constant(3.0)
with tf.Session() as sess:
    while True:
        print sess.run(a)

但是它仍然在我的两个 gpu 中分配内存。

|    0      7479    C   python                         5437MiB 
|    1      7479    C   python                         5437MiB 

Answer 1

相信你需要设置CUDA_VISIBLE_DEVICES=1 .或者你想使用哪个 GPU。如果你只让一个 GPU 可见，你将把它称为 /gpu:0在 tensorflow 中，无论您将环境变量设置为什么。

有关该环境变量的更多信息:https://devblogs.nvidia.com/cuda-pro-tip-control-gpu-visibility-cuda_visible_devices/