深度学习pipeline - 自动识别空闲GPU

# !bin/bash
# 一些英伟达的命令
nvidia-smi -h  # 帮助

nvidia-smi -q / nvidia-smi --query  # 显示所有gpu的信息.

nvidia-smi -q -d / nvidia-smi -q --display= # 只显示指定的信息: MEMORY, UTILIZATION, PIDS等.

知道这些基本的英伟达的命令之后，我们就可以用它来识别空闲的GPU了。

import os
command = os.popen("nvidia-smi -q -d PIDS | grep Processes")
lines = command.read().split("\n") # 如果显卡上有进程那么这一行只会有一个Processes
free_gpu = []
for i in range(len(lines)):
    if "None" in lines[i]:
        free_gpu.append(i)

print(free_gpu)

然后我们就可以去设置

os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(get_free_gpu())