天天看點

tensorflow gpu版本 檢視可用gpu ,指定 gpu ,禁用gpu方法

安裝gpu版本的tensorflow

pip install tensorflow-gpu 

# 預設安裝最新版 如果需要指定版本的話使用,如1.11.0
pip install tensorflow-gpu==1.11.0
           

檢視可用gpu

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())
           

我的輸出結果如下

[name: "/device:CPU:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}

, name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 1451291444
locality {
  bus_id: 1
  links {
  }
}
           

這說明我可以用gpu也可以用cpu來訓練模型。

指定gpu訓練模型,如 gpu 0

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0" #指定gpu 0訓練模型
           

看到有人說單台機器多gpu訓練的時候,如0,1号gpu訓練,上面 ="0" 改成   ="0,1",但是也有人說這樣的效果不好,沒有試過,感興趣的可以自己查一下。

禁用gpu的話

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

即指定gpu -1訓練模型,但是這個gpu是不存在的,預設從0開始,是以也就不能使用gpu了。

但如果想用gpu的話,可以把oos.environ["CUDA_VISIBLE_DEVICES"] = "-1" 這一行注釋掉即可。

繼續閱讀