【TensorFlow】使用指定的GPU

现有的模型训练中,一般config文件会指定需要使用的GPU。

但是在某些程序中,可能代码写的有些问题,即使指定了GPU,也会存在使用了所有空闲GPU的情况。

其实如果机器中有多块GPU,那么TensorFlow默认会占满所有能用的显存,而在需要运行多个项目时,我们希望可以使用特定的某块GPU。

下面是在指定无效时的一种万能强制性指定方式:

1. 在启动文件的开头加上以下几行代码:

  • TensorFlow中指定使用某块GPU

import
os os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = "1" #默认从0开始,此处即为使用第2块GPU

  • TensorFlow中指定某几块GPU

import
os os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1" #使用0和1号GPU,即第1和2块GPU

  • TensorFlow中禁用GPU

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

2. 在启动时的命令中添加:

  • 命令前添加 CUDA_VISIBLE_DEVICES

CUDA_VISIBLE_DEVICES=0 python train.py

猜你喜欢

转载自www.cnblogs.com/wlhr62/p/10763857.html