如何基于PCI总线ID选择GPU？

Question

我想知道是否可以基于pci总线ID或uuid在Tensorflow和Pytorch中选择GPU？

Answer 1

您可以通过设置CUDA_VISIBLE_DEVICES环境变量来实现。 假设您要在PCI总线ID = 1的GPU上训练模型，则需要执行类似的操作

$ CUDA_VISIBLE_DEVICES=1 my_training_code.py

您还可以通过在Python代码中使用CUDA_VISIBLE_DEVICES设置os.environ来实现相同的目的

os.environ['CUDA_VISIBLE_DEVICES'] = '1'
...

Answer 2

经过张量流测试

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="0" #Select GPU number 0