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

Question

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

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

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

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

Answer 1

Я не знаю о решении pure-TensorFlow. Проблема в том, что существующее место для конфигураций TensorFlow является конфигурацией сеанса. Однако для памяти GPU пул памяти GPU используется для всех сеансов 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