在Tensorflow和Keras机器学习环境中使用allow_growth显存选项

CYJ0go 2018-10-14

当我们拥有与多个用户共享的GPU计算机时,我们遇到了问题。大多数用户在他们的Tensorflow或Keras机器学习环境中运行他们的GPU进程而没有设置“allow_growth”选项。它会导致显卡的显存完全分配给该进程。实际上,它可能只需要一小部分显存来进行操作。

在Tensorflow和Keras机器学习环境中使用allow_growth显存选项

通过启用“allow_growth”选项的两个图形卡可以共享三个进程的示例

为了解决这个机器学习问题,只需在Tensorflow或Keras中启用“allow_growth”设置。

以下用于在Tensorflow中设置allow_growth显存选项的Python代码

# Tensorflow
import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config, ...)

在Tensorflow和Keras机器学习环境中使用allow_growth显存选项

对于Keras,Python代码如下:

#For Keras
from keras.callbacks import ModelCheckpoint
from keras.models import Model, load_model, save_model, Sequential
from keras.layers import Dense, Activation, Dropout, Input, Masking, TimeDistributed, LSTM, Conv1D
from keras.layers import GRU, Bidirectional, BatchNormalization, Reshape
from keras.optimizers import Adam
from keras.backend.tensorflow_backend import set_session
import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.allow_growth = True # dynamically grow the memory used on the GPU
config.log_device_placement = True # to log device placement (on which device the operation ran)
sess = tf.Session(config=config)
set_session(sess) # set this TensorFlow session as the default session for Keras

在Tensorflow和Keras机器学习环境中使用allow_growth显存选项

在机器学习过程中,这增加了显卡的利用率,且不限制主机所具有的卡数量。

相关推荐