我们使用以下命令指定要使用的 GPU 设备:
with tf.device('/gpu:'+gpu_id):
gpu_id是一个字符串变量,我在其中手动设置要使用的 GPU ID。我需要运行多个实验,每个实验在不同的 GPU 上。因此,我在运行代码实例之前手动更改gpu_id的值。
我可以编写一些代码来自动检测第一个未使用的 GPU 并将其设置为 gpu_id 吗?
已经有一个函数可以让你知道哪个GPU用于张量:
# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
将log_device_placement
设置为 True
将返回类似于以下内容的数据:
Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: Tesla K40c, pci bus
id: 0000:05:00.0
b: /job:localhost/replica:0/task:0/gpu:0
a: /job:localhost/replica:0/task:0/gpu:0
MatMul: /job:localhost/replica:0/task:0/gpu:0
[[ 22. 28.]
[ 49. 64.]]
↳ 使用图形用户界面
我正在研究 TF-2.1 和火炬,所以我不想在任何 ML 帧中指定这个自动选择。我只是使用原始的nvidia-smi和os.environ来获得一个空置的GPU。
def auto_gpu_selection(usage_max=0.01, mem_max=0.05):
"""Auto set CUDA_VISIBLE_DEVICES for gpu
:param mem_max: max percentage of GPU utility
:param usage_max: max percentage of GPU memory
:return:
"""
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
log = str(subprocess.check_output("nvidia-smi", shell=True)).split(r"n")[6:-1]
gpu = 0
# Maximum of GPUS, 8 is enough for most
for i in range(8):
idx = i*3 + 2
if idx > log.__len__()-1:
break
inf = log[idx].split("|")
if inf.__len__() < 3:
break
usage = int(inf[3].split("%")[0].strip())
mem_now = int(str(inf[2].split("/")[0]).strip()[:-3])
mem_all = int(str(inf[2].split("/")[1]).strip()[:-3])
# print("GPU-%d : Usage:[%d%%]" % (gpu, usage))
if usage < 100*usage_max and mem_now < mem_max*mem_all:
os.environ["CUDA_VISIBLE_EVICES"] = str(gpu)
print("nAuto choosing vacant GPU-%d : Memory:[%dMiB/%dMiB] , GPU-Util:[%d%%]n" %
(gpu, mem_now, mem_all, usage))
return
print("GPU-%d is busy: Memory:[%dMiB/%dMiB] , GPU-Util:[%d%%]" %
(gpu, mem_now, mem_all, usage))
gpu += 1
print("nNo vacant GPU, use CPU insteadn")
os.environ["CUDA_VISIBLE_EVICES"] = "-1"
如果我能得到任何 GPU,它会CUDA_VISIBLE_EVICES设置为该 GPU 的 BUSID:
GPU-0 is busy: Memory:[5738MiB/11019MiB] , GPU-Util:[60%]
GPU-1 is busy: Memory:[9688MiB/11019MiB] , GPU-Util:[78%]
Auto choosing vacant GPU-2 : Memory:[1MiB/11019MiB] , GPU-Util:[0%]
否则,设置为 -1 以使用 CPU:
GPU-0 is busy: Memory:[8900MiB/11019MiB] , GPU-Util:[95%]
GPU-1 is busy: Memory:[4674MiB/11019MiB] , GPU-Util:[35%]
GPU-2 is busy: Memory:[9784MiB/11016MiB] , GPU-Util:[74%]
No vacant GPU, use CPU instead
注意:在导入任何需要 GPU 的 ML 帧之前,请使用此功能,然后它可以自动选择 GPU。此外,您可以轻松地设置多个任务。