解决:CUDA_ERROR_OUT_OF_MEMORY 问题

2019-11-27  本文已影响0人  胜负55开

前言:在用Keras或直接Tensorflow训练大型网络时,经常会报如题目中"显存不足"不足的错误。其实绝大多数情况:只是tensorflow一个人把所有的显存都先给占了(程序默认的),导致其他需要显存的程序部分报错!

完整的处理很简单,可分下面简单的3步:

其实操作很简单,看下面的例子即可:假定直接用的Keras

import os   # 设置程序所使用的显卡(编号)
import tensorflow as tf
import keras
import keras.backend.tensorflow_backend as KTF   # 修改Tensorflow的后端

# 设置程序使用的GPU编号:我用6号和7号显卡
os.environ['CUDA_VISIBLE_DEVICES'] = '6,7'

# 修改Keras的后端:Tensorflow中的一些"不好的"默认设置:
config = tf.ConfigProto()
config.gpu_options.allow_growth = True  # 按需使用显存,而不是一次性先全占着!
config.gpu_options.per_process_gpu_memory_fraction = 0.7   # 最大显存占用率70%
config.allow_soft_placement = True  # 可以自带选择当前可用的GPU设备:6或7都能用;
config.log_device_placement = False  # 程序运行中的各种打印,是否顺便显示出于哪台设备?—— 一般不用 

sess = tf.Session(config=config)
KTF.set_session( sess )  # 设置Tensorflow的Session,同时也就设置好了Keras!

# -------------------- 下面就正常的模型搭建、编译、训练 -------------------- #
model.add( layers.xx ) ....

model.compile() ....

model.fit() ...

另外,有时batch_size的大小也会导致显存问题!一般不要太大,也不要太小。

上一篇下一篇

猜你喜欢

热点阅读