在Tensorflow 训练模型时报错提示: 

failed to allocate 3.77G (4046333952 bytes) from device: CUDA_ERROR_OUT_OF_MEMORY: out of memory

虽然会报出显存溢出问题,但不影响正常训练,不过笔者还是想知道这个问题是怎么来的。

废话不多说,先上session初始化的代码

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=1.0)

config = tf.ConfigProto(gpu_options=gpu_options, allow_soft_placement=False, log_device_placement=False)

session = tf.Session(config=config)

per_process_gpu_memory_fraction 代表显卡可用使用的占比,代码中设置为 1 ,代表显卡被100%使用,所以会在训练中报出上述错误提示

现在,修改一下session初始化代码

config = tf.ConfigProto(allow_soft_placement=False, log_device_placement=False) 

config.gpu_options.allow_growth = True 

session = tf.Session(config=config)

在tf.ConfigProto()参数中,删除gpu_options的设定,并且在初始化config后,新增 config.gpu_options.allow_growth = True ,这个选项是让session根据当前显存大小,自动分配显存,不会将显存全部占用,从而确保最大占用比。

Logo

旨在为数千万中国开发者提供一个无缝且高效的云端环境,以支持学习、使用和贡献开源项目。

更多推荐