1. 程式人生 > >tensorflow獲取當前可用的GPU資訊

tensorflow獲取當前可用的GPU資訊

簡而來講,我目前電腦中有三個可用的GPU,將返回['/ gpu:0','/ gpu:1','/ gpu:2'],可以使用以下程式碼檢查所有裝置列表:

from tensorflow.python.client import device_lib

def get_available_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_available_gpus())