TensorFlow GPU 显存分配策略：按需增长与比例限制详解

TensorFlow 默认显存占用机制

在默认配置下，TensorFlow 会尝试映射进程可见的所有 GPU 设备（受 CUDA_VISIBLE_DEVICES 环境变量影响）上的绝大部分显存。这种策略旨在通过预分配大块连续内存来减少运行时的内存碎片，从而更有效地利用设备上相对稀缺的 GPU 显存资源。

然而，在某些应用场景中，理想状态是进程仅分配可用显存的一个子集，或者根据实际运行需求动态增加显存占用。为此，TensorFlow 提供了 Session 级别的两个配置选项，允许开发者对显存分配行为进行精细控制。

动态显存分配策略

第一种策略是自主申请所需的内存空间，即允许 TensorFlow 根据运行时的实际需求来动态分配 GPU 显存。

具体实现为 allow_growth 选项。该选项在初始化时仅分配极少量的显存，随着 Session 的运行和模型对显存需求的增加，TensorFlow 会逐步扩展已分配的 GPU 内存区域。需要注意的是，一旦分配的显存被使用，系统不会主动释放，这是为了避免因频繁释放和重新分配导致的严重内存碎片问题。开启此选项需在 ConfigProto 中设置相应参数，具体配置代码如下：

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config, ...)

固定比例显存限制策略

第二种策略是规定可用最大内存占单个 GPU 总显存的比例，即通过固定比例限制显存使用。

具体实现为 per_process_gpu_memory_fraction 选项。该选项决定了每个可见 GPU 设备中，当前进程可分配到的显存占总显存量的比例。例如，若希望 TensorFlow 仅使用每个 GPU 总显存的 40%，可通过以下配置实现：

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4
session = tf.Session(config=config, ...)

上述配置方法对于真正限制 TensorFlow 进程可使用的 GPU 显存总量非常实用，尤其在多模型共存或显存受限的环境中，能有效防止显存溢出。

以上关于 TensorFlow 显存使用机制及配置方法的介绍到此结束，希望能为读者在模型训练与部署中的显存管理提供参考。