Tensorflow на общих графических процессорах: как автоматически выбрать тот, который не используется

У меня есть доступ через ssh к кластеру из n GPU. Tensorflow автоматически дал им имена gpu:0,...,gpu:(n-1).

У других тоже есть доступ, и иногда они берут случайный gpu. Я не размещал никакие tf.device() явно, потому что это громоздко, и даже если бы я выбрал номер j-й видеокарты и что кто-то уже использует номер j-й видеокарты, это было бы проблематично.

Я хотел бы просмотреть использование графического процессора и найти первый, который не используется, и использовать только этот. Я думаю, кто-то мог бы проанализировать вывод nvidia-smi с помощью bash, получить переменную i и передать эту переменную i скрипту tensorflow в качестве номера используемого графического процессора.

Я никогда не видел ни одного примера этого. Я думаю, что это довольно распространенная проблема. Каким будет самый простой способ сделать это? Доступен ли чистый тензорный поток?


tensorflow gpu distributed-system
person jeandut    schedule 13.01.2017    source источник

Ответы (2)


arrow_upward
17
arrow_downward

Я не знаю о чистом решении TensorFlow. Проблема в том, что существующим местом для конфигураций TensorFlow является конфигурация сеанса. Однако для памяти графического процессора пул памяти графического процессора является общим для всех сеансов TensorFlow внутри процесса, поэтому конфигурация сеанса — неподходящее место для его добавления, а механизма для глобальной конфигурации процесса нет (но он должен быть, чтобы также возможность настроить глобальный процесс Eigen threadpool). Поэтому вам нужно сделать это на уровне процесса, используя переменную среды CUDA_VISIBLE_DEVICES.

Что-то вроде этого:

import subprocess, re

# Nvidia-smi GPU memory parsing.
# Tested on nvidia-smi 370.23

def run_command(cmd):
    """Run command, return output as string."""
    output = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True).communicate()[0]
    return output.decode("ascii")

def list_available_gpus():
    """Returns list of available GPU ids."""
    output = run_command("nvidia-smi -L")
    # lines of the form GPU 0: TITAN X
    gpu_regex = re.compile(r"GPU (?P<gpu_id>\d+):")
    result = []
    for line in output.strip().split("\n"):
        m = gpu_regex.match(line)
        assert m, "Couldnt parse "+line
        result.append(int(m.group("gpu_id")))
    return result

def gpu_memory_map():
    """Returns map of GPU id to memory allocated on that GPU."""

    output = run_command("nvidia-smi")
    gpu_output = output[output.find("GPU Memory"):]
    # lines of the form
    # |    0      8734    C   python                                       11705MiB |
    memory_regex = re.compile(r"[|]\s+?(?P<gpu_id>\d+)\D+?(?P<pid>\d+).+[ ](?P<gpu_memory>\d+)MiB")
    rows = gpu_output.split("\n")
    result = {gpu_id: 0 for gpu_id in list_available_gpus()}
    for row in gpu_output.split("\n"):
        m = memory_regex.search(row)
        if not m:
            continue
        gpu_id = int(m.group("gpu_id"))
        gpu_memory = int(m.group("gpu_memory"))
        result[gpu_id] += gpu_memory
    return result

def pick_gpu_lowest_memory():
    """Returns GPU with the least allocated memory"""

    memory_gpu_map = [(memory, gpu_id) for (gpu_id, memory) in gpu_memory_map().items()]
    best_memory, best_gpu = sorted(memory_gpu_map)[0]
    return best_gpu

Затем вы можете поместить его в utils.py и установить GPU в своем скрипте TensorFlow перед первым tensorflow импортом. IE

import utils
import os
os.environ["CUDA_VISIBLE_DEVICES"] = str(utils.pick_gpu_lowest_memory())
import tensorflow
person Yaroslav Bulatov    schedule 13.01.2017
comment
Спасибо за этот блестящий ответ! - person jeandut; 14.01.2017
comment
Таким образом, очевидно, что nvidia-smi в некоторых ситуациях может давать несовпадающие номера устройств, похоже, вам нужно объединить их lspci, чтобы получить правильные номера, как описано в 152 - person Yaroslav Bulatov; 18.01.2017
comment
Я проверю это спасибо! Но до сих пор ваше решение, казалось, работало для меня нормально! - person jeandut; 19.01.2017
comment
Если он перестанет работать, обходной путь установит переменную среды: export CUDA_DEVICE_ORDER=PCI_BUS_ID - person Yaroslav Bulatov; 20.01.2017
comment
Я видел это на гитхабе. Спасибо еще раз ! - person jeandut; 20.01.2017

arrow_upward
8
arrow_downward

Реализация решения Ярослава Булатова доступна на https://github.com/bamos/setGPU. .

person Trisoloriansunscreen    schedule 27.04.2018