使用gpu版的tensorflow时出现OOM错误(显存溢出错误)： CUDA out of memory的解决办法。

无码不欢的我

已于 2022-03-17 10:05:52 修改

阅读量5.3k

点赞数

分类专栏： bug信息解决方案文章标签： tensorflow 深度学习人工智能

于 2022-03-15 09:39:16 首次发布

本文链接：https://blog.csdn.net/m0_56192771/article/details/123494264

版权

bug信息解决方案专栏收录该内容

10 篇文章

订阅专栏

本文介绍了如何通过设置`tf.config`来限制TensorFlow在GPU上的内存使用，避免程序因内存不足而报错，特别是在显存为6GB的设备上，建议将内存限制设置为4GB。遇到'RuntimeError: Virtual devices cannot be modified after being initialized'时，需重启电脑并重新配置。

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

gpu版的tensorflow 在执行过程中会默认使用全部的 GPU 内存，给系统保留 200 M，这样容易造成显存溢出，因此我们可以使用如下语句指定 GPU 内存的分配空间：

下面以我自己的电脑为例，进行相关设置。因为我的电脑的显存是6G，所以我可以设置tensorflow程序运行时所占用的显存不能超过6G，这里限制在4G大小。

import tensorflow as tf
#限制消耗固定大小的显存（程序不会超出限定的显存大小，若超出的报错）。
physical_gpus = tf.config.list_physical_devices("GPU")
tf.config.experimental.set_virtual_device_configuration(
    physical_gpus[0], #对0号gpu进行限制
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4000)] 
)
logical_gpus = tf.config.list_logical_devices("GPU")
# 注：memory_limit是限制显存大小，我显卡6G，这里设置的4000MB

另外需要特别注意的是，如果在执行上述代码时出现如下报错信息：
RuntimeError: Virtual devices cannot be modified after being initialized
则需要重新启动一下电脑，然后重新执行上述代码就可以了。