Allowing GPU memory growth

最新推荐文章于 2023-09-26 17:40:18 发布

Microstrong0305

最新推荐文章于 2023-09-26 17:40:18 发布

阅读量2.6k

点赞数 1

分类专栏：深度学习 TensorFlow学习笔记文章标签： Tensorflow GPU

深度学习同时被 2 个专栏收录

34 篇文章

订阅专栏

TensorFlow学习笔记

20 篇文章

订阅专栏

本文介绍如何在TensorFlow中管理和控制GPU内存使用。包括使用allow_growth选项按需分配GPU内存，以及通过per_process_gpu_memory_fraction选项限定进程可使用的GPU内存量。

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

默认情况下，TensorFlow会占用所有GPUs的所有GPU内存（取决于CUDA_VISIBLE_DEVICES这个系统变量），这样做可以减少内存碎片，更有效地利用设备上相对宝贵的GPU内存资源。

在某些情况下，该进程仅仅需要分配可用内存的一部分，或者根据该进程的需要来增加内存的使用量。TensorFlow在Session上提供了两个Config选项来进行控制。

第一个是“allow_growth”选项，它仅仅基于运行时的分配来分配更多的GPU内存：它开始分配非常少的内存，并且随着Session的运行和更多的GPU内存需求，扩展TensorFlow所需的GPU内存区域。这可能导致很糟糕的内存碎片。要打开此选项，请在ConfigProto中将设置为：

#allow growth
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config, ...)
# 使用allow_growth option，刚一开始分配少量的GPU容量，然后按需慢慢的增加，由于不会释放内存，所以会导致碎片

第二个方法是“pre_process_gpu_memory_fraction“选项，它决定了每个可见的GPU应分配的内存总量的百分比。例如，你可以告诉TensorFlow仅仅分配总内存的40%，通过以下设定就可以实现：

# per_process_gpu_memory_fraction
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4
session = tf.Session(config=config, ...)
#设置每个GPU应该拿出多少容量给进程使用，0.4代表 40%

如果你想限制TensorFlow进程可利用的GPU内存数量，以上对你非常的有用。

Reference：

【1】https://www.tensorflow.org/guide/using_gpu

【2】TensorFlow炼丹（1） Using GPUs - 张天亮的文章 - 知乎 https://zhuanlan.zhihu.com/p/28083241

【3】https://blog.youkuaiyun.com/u012436149/article/details/53837651