深度學習模型指定GPU訓練

Pytorch

# 1：

torch.cuda.set_device(1)

# 2：

device = torch.device("cuda:1")

# 3：（官方推薦）

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

# 4：（同時調用兩塊GPU的話） 

os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'

Keras

注意：先安裝tensorflow-gpu版本，再安裝keras，這樣keras才能使用GPU加速。

# -------------------------------------指定GPU運行--------------------------------------
# os.environ["CUDA_VISIBLE_DEVICES"]="0"
# -------------------------------------指定GPU運行--------------------------------------
# -------------------------------------指定CPU運行--------------------------------------
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
# -------------------------------------指定CPU運行--------------------------------------

TensorFlow