如何在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) 

但它仍然在我的兩個gpus中分配內存。

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

Answer 1

我相信你需要設置CUDA_VISIBLE_DEVICES=1。或者您想要使用哪種GPU。如果只顯示一個GPU，則無論您將環境變量設置爲什麼，都會將其稱爲/gpu:0。

對環境變量的更多信息：http://www.acceleware.com/blog/cudavisibledevices-masking-gpus

Answer 2

有3種方法來實現這一目標：

1. 使用CUDA_VISIBLE_DEVICES環境變量。 通過設置設置環境變量CUDA_VISIBLE_DEVICES=1僅使設備1可見，並通過設置CUDA_VISIBLE_DEVICES=0,1使設備0和1可見。使用with tf.device('/gpu:2')並創建圖形。然後它將使用GPU設備2運行。使用config = tf.ConfigProto(device_count = {'GPU': 1})然後sess = tf.Session(config=config)。這將使用GPU設備1.

Answer 3

如果沒有告知其他情況，TF將分配每個可見GPU上的所有可用內存。

如果是一次性事物，則可以設置CUDA_VISIBLE_DEVICES=0,1。

如果您正在使用羣集，並且不想跟蹤哪個GPU處於繁忙狀態並手動鍵入信息，則可在構建會話之前調用下一個方法。它會過濾掉已經在使用的GPU（沒有足夠的可用內存）併爲您設置CUDA_VISIBLE_DEVICES。

功能：

import subprocess as sp 
import os 

def mask_unused_gpus(leave_unmasked=1): 
    ACCEPTABLE_AVAILABLE_MEMORY = 1024 
    COMMAND = "nvidia-smi --query-gpu=memory.free --format=csv" 

    try: 
    _output_to_list = lambda x: x.decode('ascii').split('\n')[:-1] 
    memory_free_info = _output_to_list(sp.check_output(COMMAND.split()))[1:] 
    memory_free_values = [int(x.split()[0]) for i, x in enumerate(memory_free_info)] 
    available_gpus = [i for i, x in enumerate(memory_free_values) if x > ACCEPTABLE_AVAILABLE_MEMORY] 

    if len(available_gpus) < leave_unmasked: raise ValueError('Found only %d usable GPUs in the system' % len(available_gpus)) 
    os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, available_gpus[:leave_unmasked])) 
    except Exception as e: 
    print('"nvidia-smi" is probably not installed. GPUs are not masked', e) 

限制：如果你開始一次它仍然會引起衝突，因爲當你創建一個會話內存沒有immideately分配多個腳本。如果它是你有問題，你可以使用一個隨機版本，在我原來的source code: mask_busy_gpus()