第1章：TensorFlow显存管理机制

大家好，我是你们的讲师。今天咱们聊聊TensorFlow的显存管理。说实话，我见过太多同学在GPU集群上跑模型，一上来就OOM（显存溢出），然后一脸懵地问我：“老师，我代码明明没问题啊！”

嗯，问题往往出在显存管理上。TensorFlow和PyTorch在这块儿差别挺大的。我习惯把显存管理比作“租房”——你租多大房子、怎么分配房间、房东（框架）怎么管，都直接影响你能不能住得舒服。

1.1 GPU内存增长配置

先说说最基础的东西——GPU内存增长。TensorFlow默认会“霸占”整张显卡的显存。你想想看，这多浪费？我跑个小模型，结果它把12GB显存全占了，别的任务只能干瞪眼。

⚠️ 注意： 默认情况下，TensorFlow会一次性申请全部显存。这在多任务场景下非常不友好。

怎么解决？用 tf.config.experimental.set_memory_growth。我个人习惯在代码开头就加上这个配置：

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)

这段代码的意思是：让TensorFlow按需申请显存，而不是一口气全占满。我在项目中遇到过，有个同事没加这个配置，结果在8卡机器上跑单卡任务，其他7张卡全被“冻住”了。加了之后，显存占用从12GB降到了2GB左右。

💡 小技巧： 如果你想让TensorFlow只使用特定比例的显存，可以用 tf.config.set_logical_device_configuration 来设置上限。比如限制为4GB：

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

1.2 显存分配策略

TensorFlow的显存分配策略，说白了就是“怎么分蛋糕”。它有两种模式：

贪婪分配（Greedy Allocation）：默认模式，一次性申请全部显存。简单粗暴，但浪费。
动态分配（Dynamic Allocation）：按需增长，用多少申请多少。更灵活，但可能有碎片问题。

我建议在训练大模型时用动态分配，推理时用固定分配。为什么？训练时显存需求波动大，动态分配能避免浪费；推理时显存需求稳定，固定分配能减少碎片。

这里有个坑——我曾经在推理服务中用了动态分配，结果发现显存碎片越来越多，最后OOM了。排查了半天，发现是频繁创建和销毁Tensor导致的。解决办法是：用 tf.function 把推理过程编译成图，减少临时变量。

@tf.function
def inference_step(inputs):
    # 推理逻辑
    return model(inputs)

# 这样显存分配会更稳定

1.3 TF与PyTorch的差异对比

很多同学问我：“老师，TF和PyTorch到底哪个显存管理好？”我的回答是：各有千秋，看场景。

咱们直接上对比表：

对比维度	TensorFlow	PyTorch
默认显存策略	全占（贪婪分配）	按需分配（动态）
显存增长配置	需手动开启 `set_memory_growth`	默认就是动态增长
显存释放	依赖图执行，释放较慢	引用计数，释放及时
多卡支持	`tf.distribute.MirroredStrategy`	`torch.nn.DataParallel` / `DistributedDataParallel`
显存碎片问题	较严重（图执行导致）	较轻（动态图更灵活）
调试工具	`tf.profiler`、`tf.debugging`	`torch.cuda.memory_summary()`

从表里能看出来，PyTorch在显存管理上更“亲民”。但TensorFlow在大型分布式训练中，图执行的优势很明显——它能做全局优化，减少显存峰值。

我记得有一次，用TF跑一个BERT-large模型，显存峰值是11.2GB；换成PyTorch同样的配置，峰值到了13.5GB。为什么？因为TF的图执行能复用中间变量，而PyTorch的动态图会保留更多临时数据。

📌 核心结论：
小模型、快速实验 → 选PyTorch（显存管理省心）
大模型、生产环境 → 选TensorFlow（图优化更省显存）
混合使用？可以！用TF训练，用PyTorch推理（我见过不少团队这么干）

知识体系总览

下面这张图，是我自己总结的TF显存管理知识体系。你把它存下来，以后遇到OOM问题，按图索骥就行。

这张图把TF显存管理的三个核心维度串起来了。你从“内存增长配置”开始，再到“分配策略”，最后对比“TF vs PyTorch”，基本就能掌握80%的显存管理知识。

最后说一句：别怕OOM。我刚开始做分布式训练时，一天能遇到七八次OOM。后来慢慢摸清了这些机制，现在基本能提前预判显存瓶颈。你也能做到。