3、TensorFlow显存管理：分配策略、API详解与碎片化问题

显存管理，说白了就是怎么让GPU这块“金贵”的内存别浪费。TensorFlow在这块做得挺灵活，但也容易踩坑。我刚开始用TF的时候，经常遇到“OOM”报错，后来才发现是分配策略没搞对。

3.1 TensorFlow的显存分配策略

TensorFlow默认会“贪心”地占满所有可见的GPU显存。你想想看，一个模型可能只用2GB，它却把8GB全占了。这其实是为了避免运行时动态分配带来的性能开销。但坏处也很明显——别的程序就没法用了。

我个人习惯是，除非在做推理服务，否则很少用默认策略。为什么呢？因为调试的时候，你往往需要同时跑好几个实验。显存全占满了，就只能排队等了。

TF提供了两种主流策略：

• 预分配全部显存：默认行为，启动时就申请所有显存。
• 按需增长：用多少申请多少，可以设置上限。

嗯，这里要注意：按需增长虽然灵活，但频繁申请释放会引入额外延迟。如果你的模型大小是固定的，预分配反而更快。

3.2 tf.config.gpu API详解

这个API是控制显存的核心工具。我建议你把它当成“显存遥控器”来用。下面这几个方法，我几乎每个项目都会用到。

3.2.1 设置显存按需增长

import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        print("显存按需增长已开启")
    except RuntimeError as e:
        print(e)

这段代码我几乎每个项目都会贴一遍。记得要在创建任何张量或模型之前调用。我曾经因为顺序搞反了，调试了半天才发现是API调用时机不对。

3.2.2 限制显存使用上限

有时候你只想让TF用一部分显存，比如留2GB给数据加载或别的进程。这时候可以用虚拟设备限制：

gpus = tf.config.list_physical_devices('GPU')
if gpus:
    tf.config.set_logical_device_configuration(
        gpus[0],
        [tf.config.LogicalDeviceConfiguration(memory_limit=4096)]
    )

这里设了4GB上限。我有个项目需要同时跑两个模型，一个占3GB，一个占2GB。用这个API就能精确分配，互不干扰。

3.2.3 查看显存使用情况

print(tf.config.experimental.get_memory_info('GPU:0'))

返回一个字典，包含当前使用量和峰值使用量。调试显存泄漏时特别好用。

3.3 显存碎片化问题

这是个比较隐蔽的问题。你可能会遇到这种情况：明明显存总量够，但一训练就报OOM。为什么会这样？

说白了，就是显存被“切碎”了。比如你反复创建和销毁不同大小的张量，显存里就会出现很多小空洞。每个空洞单独看都不大，但加起来很可观。新的大张量需要连续空间时，就申请失败了。

我在项目中遇到过最夸张的一次，8GB显存只用了4GB，但就是跑不动一个3GB的模型。查了半天，发现是数据预处理时频繁创建临时张量导致的。

3.3.1 如何避免碎片化

• 预分配大张量：尽量复用已有的张量，而不是反复创建。
• 使用tf.function：将计算图编译成静态图，减少动态分配。
• 调整batch size：有时候减小batch size反而能跑得更快，因为避免了碎片化导致的OOM。

3.3.2 碎片化的可视化

你可以用nvidia-smi或者TensorBoard的显存分析工具来观察。不过说实话，最直接的办法还是用上面提到的get_memory_info来监控峰值。

3.4 知识体系图

下面这张图帮你理清本章的核心逻辑：

这张图展示了三个核心模块的关联。分配策略决定了TF怎么拿显存，API给了你精细控制的能力，而碎片化则是实际运行中容易踩的坑。三者环环相扣。

3.5 实战建议

最后，给你几个我常用的配置模板：

场景	推荐策略	API设置
单模型训练	预分配全部	默认即可
多模型并行	按需增长 + 上限	set_memory_growth + memory_limit
推理服务	预分配 + 固定batch	memory_limit 精确设置
调试阶段	按需增长	set_memory_growth

嗯，显存管理这块，说白了就是“预则立，不预则废”。花几分钟配置好，能省下几小时的调试时间。我个人习惯是，每个新项目的第一件事就是写好显存配置代码，然后再开始搭模型。