keras在使用GPU的时候有个特点,就是默认全部占满显存。
若单核GPU也无所谓,若是服务器GPU较多,性能较好,全部占满就太浪费了。
于是乎有以下五种情况:
本节来源于:深度学习theano/tensorflow多显卡多人使用问题集(参见:Limit the resource usage for tensorflow backend · Issue #1538 · fchollet/keras · GitHub)
在使用keras时候会出现总是占满GPU显存的情况,可以通过重设backend的GPU占用情况来进行调节。
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.3
set_session(tf.Session(config=config))
需要注意的是,虽然代码或配置层面设置了对显存占用百分比阈值,但在实际运行中如果达到了这个阈值,程序有需要的话还是会突破这个阈值。换而言之如果跑在一个大数据集上还是会用到更多的显存。以上的显存限制仅仅为了在跑小数据集时避免对显存的浪费而已。(2017年2月20日补充)
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
此时的代码为选择了编号为2 的GPU
# python设置系统变量的方法
os.environ["CUDA_VISIBLE_DEVICES"] = "8,9,10,11,12,13,14,15"
注意,在代码中指定设备时,重新从0开始计,而不是从8开始。
来源:Tensorflow 学习笔记(七) ———— 多GPU操作
上述两个连一起用就行:
import os
import tensorflow as tf
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.3
set_session(tf.Session(config=config))
那么在命令行,可以使用:https://github.com/tensorflow/nmt/issues/60
CUDA_VISIBLE_DEVICES=0 python -m nmt.nmt
import keras.backend.tensorflow_backend as KTF
import tensorflow as tf
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
config = tf.ConfigProto()
config.gpu_options.allow_growth=True #不全部占满显存, 按需分配
sess = tf.Session(config=config)
KTF.set_session(sess)
os.environ指的是占用的GPU编号;allow_growth为动态申请显存占用。
来自博客:TensorFlow如何充分使用所有CPU核数,提高TensorFlow的CPU使用率,以及Intel的MKL加速
num_cores = 4
config = tf.ConfigProto(intra_op_parallelism_threads=num_cores, inter_op_parallelism_threads=num_cores,
allow_soft_placement=True, device_count={'CPU': 4})
session = tf.Session(config=config)
K.set_session(session)
其中:
DistributionStrategy API是构建多设备/机器训练的简单方式,开发者只需要在现有模型上做少量的修改,就可以用它们进行分布式训练。另外,DistributionStrategy在设计时考虑了同时兼容动态图(eager)和静态图。
参考:TensorFlow 1.11.0发布,一键多GPU(训练、预测和评价tf.keras模型)
目前TensorFlow支持三种DistributionStrategy:
在tf.keras中直接使用DistributionStrategy
最新的TensorFlow Github中给出了在tf.keras中直接使用DistributionStrategy的例子。
用tf.keras构建一个单层网络:
inputs = tf.keras.layers.Input(shape=(1,))
predictions = tf.keras.layers.Dense(1)(inputs)
model = tf.keras.models.Model(inputs=inputs, outputs=predictions)
目前,使用DistributionStrategy需要使用tf.data.Dataset来作为数据输入:
features = tf.data.Dataset.from_tensors([1.]).repeat(10000).batch(10)
labels = tf.data.Dataset.from_tensors([1.]).repeat(10000).batch(10)
train_dataset = tf.data.Dataset.zip((features, labels))
这里我们为模型指定使用MirroredStrategy进行多GPU训练,代码非常简单:
distribution = tf.contrib.distribute.MirroredStrategy()
model.compile(loss='mean_squared_error',
optimizer=tf.train.GradientDescentOptimizer(learning_rate=0.2),
distribute=distribution)
使用常规的训练、评价和预测方法会自动在多GPU上进行:
model.fit(train_dataset, epochs=5, steps_per_epoch=10)
model.evaluate(eval_dataset)
model.predict(predict_dataset)
将tf.keras模型迁移到多GPU上运行只需要上面这些代码,它会自动切分输入、在每个设备(GPU)上复制层和变量、合并和更新梯度。