Keras/pytorch在训练时指定gpu以及多gpu训练

首先在终端输入命令nvidia-smi查看gpu运行状态

可见在服务器存在两台1080ti显卡，编号分别为0,1。在实际情况中选择一块无人使用的显卡进行训练最为合适。另外在训练过程中TensorFlow和Keras都是默认占满全部显存，多模型在一张卡训练时容易造成爆显存的情况，所以采用限制显存大小的方法。

import tensorflow as tf
import keras.backend.tensorflow_backend as KTF

config = tf.ConfigProto()
# 指定可见显卡
config.gpu_options.visible_device_list="1"
#不满显存, 自适应分配
config.gpu_options.allow_growth=True   
sess = tf.Session(config=config)
KTF.set_session(sess)

若要使用多块gpu训练，使用keras.utils.multi_gpu_model方法

在原来的model基础上使用multi_gpu_model函数指定一下gpu个数：

model =  multi_gpu_model(model, 2)

pytorch则是

if torch.cuda.device_count() > 1:
        model = nn.DataParallel(model)
model.to(device)