tensorflow指定gpu和cpu

如果机器中有多块GPU,tensorflow会默认吃掉所有能用的显存,
如果实验室多人公用一台服务器,希望指定使用特定某块GPU。

可以在文件开头加入如下代码:

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # 照PCI_BUS_ID顺序从0开始排列GPU设备
os.environ["CUDA_VISIBLE_DEVICES"] = "1"       # 使用第二块GPU(从0开始)

也可以制定使用某几块GPU

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 2"    # 使用第一, 三块GPU

禁用GPU

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

如果写

os.environ["CUDA_DEVICE_ORDER"] = "1" # 如果只有一块gpu,将会用cpu
 

猜你喜欢

转载自blog.csdn.net/xd_wjc/article/details/83591314