安裝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" 這一行注釋掉即可。