一、问题:
最近在利用keras训练模型的时候出现以下错误,开始一看以为是cuDNN的版本不对,更换后还是没有效果,后面在网上查找后发现是显存的问题,更改为动态分配内存就可以了。
二、解决方法:
在训练的脚本开头添加以下代码进行动态分配内存,然后重新运行就可以。
import tensorflow as tf import keras config = tf.ConfigProto() config.gpu_options.allow_growth = True keras.backend.tensorflow_backend.set_session(tf.Session(config=config))
三、运行环境
- keras: 2.2.4
- tensorflow-gpu: 1.14.0
- CUDA: 10.0.130
- cudnn: 7.6.3