【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

 

3. 若在config文件中已指定了需要使用的GPU_LIST（TRAIN.GPU_LIST = '1'）:

• 在启动文件中添加 CUDA_VISIBLE_DEVICES

import os

os.environ['CUDA_VISIBLE_DEVICES'] = cfg.TRAIN.GPU_LIST

 

转载于:https://www.cnblogs.com/wlhr62/p/10763857.html