如何在tensorflow中设置特定的gpu? [英] How to set specific gpu in tensorflow?
问题描述
我想指定 gpu 来运行我的进程.我设置如下:
I want to specify the gpu to run my process. And I set it as follows:
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 中分配内存.
However it still allocate memory in both my two gpus.
| 0 7479 C python 5437MiB
| 1 7479 C python 5437MiB
推荐答案
我相信你需要设置CUDA_VISIBLE_DEVICES=1
.或者你想使用哪个 GPU.如果你只让一个 GPU 可见,你会在 tensorflow 中将其称为 /gpu:0
,而不管你将环境变量设置为什么.
I believe that you need to set CUDA_VISIBLE_DEVICES=1
. Or which ever GPU you want to use. If you make only one GPU visible, you will refer to it as /gpu:0
in tensorflow regardless of what you set the environment variable to.
有关该环境变量的更多信息:https://devblogs.nvidia.com/cuda-pro-tip-control-gpu-visibility-cuda_visible_devices/
More info on that environment variable: https://devblogs.nvidia.com/cuda-pro-tip-control-gpu-visibility-cuda_visible_devices/
这篇关于如何在tensorflow中设置特定的gpu?的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!