强制使用GPU

Tensorflow

对于Tensorflow，声明Session的时候加入device_count={‘gpu’:0}即可

import tensorflow as tf  

sess = tf.Session(config=tf.ConfigProto(device_count={'gpu':0}))

Keras

对于Keras,则调用后端函数，设置其使用如上定义的Session即可

import tensorflow as tf
import keras.backend.tensorflow_backend as KTF
 
KTF.set_session(tf.Session(config=tf.ConfigProto(device_count={'gpu':0})))

直接在终端运行时加入相关语句实现指定GPU的使用

export CUDA_VISIBLE_DEVICES=0 python test.py
# 运行test.py文件时，使用编号为0的GPU卡
export CUDA_VISIBLE_DEVICES=0,2 python test.py
# 运行test.py文件时，使用编号为0和2的GPU卡

在程序中使用os

import os
# 使用第一张与第三张GPU卡
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 2"