深度学习模型指定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