TensorFlow2.0 CUDA_ERROR_OUT_OF_MEMORY out of memory 解决办法

最新推荐文章于 2025-03-19 14:57:56 发布

OUTMAN80

最新推荐文章于 2025-03-19 14:57:56 发布

阅读量3.1k

点赞数 7

分类专栏：神经网络文章标签： tensorflow

本文链接：https://blog.csdn.net/OUTMAN80/article/details/106854001

版权

神经网络专栏收录该内容

3 篇文章

订阅专栏

本文介绍了一种解决在使用TensorFlow2.0的GPU版本进行训练时遇到的内存溢出问题的方法。不同于TensorFlow1.0的解决方案，本文提供了适用于TensorFlow2.0的具体代码实现，确保GPU内存的动态增长，避免训练过程中出现的CUDA错误。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

今天头一次使用TensorFlow2.0的GPU版本进行训练，训练中遇到【CUDA_ERROR_OUT_OF_MEMORY out of memory】，于是百度各种解决办法，基本能查到的都是TensorFlow1.0的用法，但此方法在TensorFlow2.0中不适用。于是经过多方查找，终于找到设置方法，在此进行记录，一是方便日后回看，二是弥补百度中各种查不到。

TensorFlow1.0用法

from keras import backend as K
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
K.set_session(sess)

TensorFlow2.0用法

physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
    for k in range(len(physical_devices)):
        tf.config.experimental.set_memory_growth(physical_devices[k], True)
        tf.config.experimental.
        print('memory growth:', tf.config.experimental.get_memory_growth(physical_devices[k]))
else:
    print("Not enough GPU hardware devices available")