使用gpu训练神经网络时,tensorflow默认会申请分配所有的gpu资源,这会导致明明可以跑的模型却无法跑起来,报OOM而无法创建sessions。所以可以做如下配置,申请动态分配GPU。
#指定使用那块GUP训练
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
config = tf.ConfigProto()
# 设置最大占有GPU不超过显存的70%
config.gpu_options.per_process_gpu_memory_fraction = 0.7
# 重点:设置动态分配GPU
config.gpu_options.allow_growth = True
# 创建session时
with tf.Session(config=config) as sess:
···
网友评论