共有GPUのTensorflow：使用されていないものを自動的に選択する方法

Question

sshを介してn個のGPUのクラスタにアクセスできます。 Tensorflowは自動的にgpu：0、...、gpu：（n-1）という名前を付けました。

他にもアクセス権があり、時にはランダムなgpusをとることもあります。 tf.device()を明示的に配置していないのは、gpu番号jを選択して誰かがgpu番号jに問題があることがあっても、それは面倒であり、

私はgpusの使用方法を知り、最初に使用されていないものを見つけ出し、この1つだけを使用したいと考えています。 誰かがnvidia-smiの出力をbashで解析して変数iを取得し、その変数iを使用するgpuの番号としてtensorflowスクリプトに送ることができると思います。

私はこれの例を見たことがありません。私はそれがかなり一般的な問題だと思います。それを行う最も簡単な方法は何でしょうか？純粋なテンソルフローが利用可能ですか？

Answer 1

私は純粋なTensorFlowソリューションを認識していません。問題は、TensorFlow構成の既存の場所がセッション構成であることです。しかし、GPUメモリの場合、GPUメモリプールはプロセス内のすべてのTensorFlowセッションで共有されるため、Session configは間違った場所に追加され、プロセス全体の設定のメカニズムはありません（ただし、プロセス全体の固有スレッドプールを構成することができます）。したがって、環境変数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に入れて最初tensorflowインポートする前に、あなたのTensorFlowスクリプトでGPUを設定することができます。 IE

import utils 
import os 
os.environ["CUDA_VISIBLE_DEVICES"] = str(utils.pick_gpu_lowest_memory()) 
import tensorflow