使用tensorflow查询GPU设备列表

from tensorflow.python.client import device_lib

def get_all_gpus():
    local_device_protos = device_lib.list_local_devices()
    return [x.name for x in local_device_protos if x.device_type == 'GPU']

print(get_all_gpus())

注意:docker容器开启privileged参数后可以访问所有的GPU设备,不管是否传入Devices参数,建议多租户隔离的云服务都禁用privileged功能。

Pengfei Ni ©2021