4. TensorFlow显存管理：TF GPU配置、显存增长策略、显存碎片化

说到TensorFlow的显存管理，我见过太多人一上来就踩坑。默认情况下，TF会一口气把GPU显存全占满——哪怕你只训练一个小模型。这其实挺坑的，尤其当你只有一块卡，还想同时跑点别的实验时。

嗯，今天我们就来聊聊怎么管好TF的显存。说白了就三件事：怎么配置GPU、怎么让显存按需增长、怎么对付碎片化。

4.1 TF GPU配置：从入门到实战

先说说最基本的。TF默认会调用所有可见的GPU，并且把显存全占上。我个人习惯，第一件事就是限制GPU可见性。

举个例子，你机器上有4张卡，但只想用第2张和第3张：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1,2'  # 只暴露GPU 1和2

import tensorflow as tf
print("可见GPU:", tf.config.list_physical_devices('GPU'))

这里有个细节要注意：CUDA_VISIBLE_DEVICES 的编号是从0开始的，但它是物理卡号。我在项目中遇到过有人写成了 '2,3' 结果发现程序跑在了第3和第4张卡上——嗯，物理卡号和逻辑卡号容易搞混。

如果你不想改环境变量，也可以用TF的API来指定：

gpus = tf.config.list_physical_devices('GPU')
tf.config.set_visible_devices(gpus[1:3], 'GPU')  # 只暴露第2、3张卡

4.2 显存增长策略：别一上来就吃满

默认情况下，TF会调用 tf.config.experimental.set_memory_growth 吗？不会。它直接申请全部显存。这其实是个历史遗留问题，早期TF为了性能，觉得一次性拿完最省事。

但现实是，你很可能需要同时跑多个模型，或者留点显存给数据预处理。这时候就需要显存增长了。

说白了就是：用多少，拿多少，不够了再申请。

gpus = tf.config.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

这段代码一开，TF就不会一开始就占满显存了。它会按需分配，比如你的模型需要2GB，它就只拿2GB。我自己的经验是，这个策略对大多数场景都够用，尤其是调试阶段。

但有个坑：显存增长一旦开启，就不能中途关闭。而且如果你在同一个进程中先创建了小模型，再创建大模型，TF可能会因为显存碎片化而申请失败。嗯，这个我们后面细说。

另一种策略是直接限制最大显存使用量：

gpus = tf.config.list_physical_devices('GPU')
tf.config.set_logical_device_configuration(
    gpus[0],
    [tf.config.LogicalDeviceConfiguration(memory_limit=4096)]  # 限制4GB
)

这个方式更粗暴，适合你知道模型大概吃多少的情况。我在做推理服务时常用这个，给每个模型分配固定大小的显存，避免互相干扰。

4.3 显存碎片化：看不见的杀手

显存碎片化，说白了就是显存被切成了一堆小块，虽然总量够，但找不到一块连续的大空间给新模型用。这问题在长时间训练或频繁创建/销毁模型时特别明显。

为什么会这样？TF在分配显存时，会按需申请一些固定大小的块。如果你先申请了2GB，释放后又申请3GB，中间的空隙可能就浪费了。久而久之，显存就像一块被啃过的蛋糕，到处都是洞。

我曾经遇到过一个案例：训练一个BERT模型，跑了200个epoch后突然OOM了。检查显存总量还有剩余，但就是分配失败。最后发现是碎片化导致的。

怎么解决？我总结了几个实用方法：

• 使用显存增长策略：set_memory_growth 可以减少一次性大块申请，降低碎片概率。
• 固定显存上限：用 memory_limit 给每个进程分配固定大小，避免动态申请带来的碎片。
• 重启进程：如果碎片化严重，最直接的办法就是重启训练脚本。虽然粗暴，但有效。
• 使用TF的内存优化器：TF 2.x 引入了 tf.config.experimental.enable_tensor_float_32_execution 等优化，可以减少显存占用。

另外，我建议你在训练循环中定期打印显存使用情况：

import tensorflow as tf

def print_memory_status():
    gpu = tf.config.list_physical_devices('GPU')[0]
    details = tf.config.experimental.get_memory_info(gpu)
    print(f"当前使用: {details['current'] / 1024**3:.2f} GB")
    print(f"峰值使用: {details['peak'] / 1024**3:.2f} GB")

这样你就能及时发现显存是否在持续增长，或者有没有异常波动。

4.4 知识体系总览

下面这张图，是我梳理的TF显存管理核心逻辑。你可以把它当成一个检查清单：

这张图把三个模块串起来了。你从左边开始，先配好GPU可见性，然后决定用增长还是固定策略，最后时刻警惕碎片化。每一步都有对应的API和最佳实践。

好了，关于TF显存管理就聊这么多。记住：别让TF默认行为坑了你，主动管理才是正道。