Keras中指定显卡或者显存使用以及强制使用CPU的方法

     在模型训练以及模型推断中，我们经常需要使用指定编号的显卡或者限制显存的使用大小，亦或不用GPU，只用CPU。具体方法如下：

1.指定使用一块或多块GPU：

import os
os.environ["CUDA_VISIBLE_DEVICES"]="0,2" #指定第1,3块GPU

2.指定显卡并控制显存使用量：

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

# 指定哪些GPU可用
os.environ["CUDA_VISIBLE_DEVICES"] = "1,4"

#进行配置，使用60%的GPU
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.6
session = tf.Session(config=config)

# 设置session
KTF.set_session(session )

3.动态按需分配显存：

与前面的方法相比，只改了一行代码（config.gpu_options.allow_growth = True）

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

# 指定哪些GPU可用
os.environ["CUDA_VISIBLE_DEVICES"] = "1,4"

#进行配置，使用60%的GPU
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config)

# 设置session
KTF.set_session(session )

4.如果我只想测试CPU运行速度，而不需要用GPU的话，代码如下：

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

参考资料：

1.https://zhuanlan.zhihu.com/p/65218239

2.https://www.cnblogs.com/yqs-0705/p/10749113.html