Keras指定运行时显卡及限制GPU用量

本文转载:keras指定运行时显卡及限制GPU用量

keras在使用GPU的时候有个特点,就是默认全部占满显存。 这样如果有多个模型都需要使用GPU跑的话,那么限制是很大的,而且对于GPU也是一种浪费。因此在使用keras时需要有意识的设置运行时使用那块显卡,需要使用多少容量。

这方面的设置一般有三种情况:

  1. 指定显卡
  2. 限制GPU用量
  3. 即指定显卡又限制GPU用量

查看GPU使用情况语句(linux)

1
2
# 1秒钟刷新一次
watch -n 1 nvidia-smi

指定显卡

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

这里指定了使用编号为2的GPU,大家可以根据需要和实际情况来指定使用的GPU。

限制GPU用量

设置使用GPU的百分比

1
2
3
4
5
6
7
8
9
10
import tensorflow as tf
import keras.backend.tensorflow_backend as KTF

#进行配置,使用30%的GPU
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.3
session = tf.Session(config=config)

# 设置session
KTF.set_session(session)

需要注意的是,虽然代码或配置层面设置了对显存占用百分比阈值,但在实际运行中如果达到了这个阈值,程序有需要的话还是会突破这个阈值。换而言之如果跑在一个大数据集上还是会用到更多的显存。以上的显存限制仅仅为了在跑小数据集时避免对显存的浪费而已。

GPU按需使用

1
2
3
4
5
6
7
8
9
import tensorflow as tf
import keras.backend.tensorflow_backend as KTF

config = tf.ConfigProto()
config.gpu_options.allow_growth=True #不全部占满显存, 按需分配
session = tf.Session(config=config)

# 设置session
KTF.set_session(sess)

指定GPU并且限制GPU用量

这个比较简单,就是讲上面两种情况连上即可

1
2
3
4
5
6
7
8
9
10
11
12
import os
import tensorflow as tf
import keras.backend.tensorflow_backend as KTF

# 指定第一块GPU可用
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

config = tf.ConfigProto()
config.gpu_options.allow_growth=True #不全部占满显存, 按需分配
sess = tf.Session(config=config)

KTF.set_session(sess)

一分一毛,也是心意。