4. TensorFlow显存管理：TF显存配置、动态增长策略、内存日志分析

说到TensorFlow的显存管理，我印象最深的是刚入行那会儿。一个模型怎么都跑不起来，报错信息就四个字——"OOM"。当时我以为是模型太大，后来才发现，是TF默认把整张显卡的显存全占了。嗯，这事儿其实挺坑的。

TF的显存管理，说白了就三个核心问题：怎么分配、怎么增长、怎么排查。我们一个一个说。

4.1 TF的显存分配策略

TF默认的行为是什么？你跑一个session，它就把GPU上所有可见的显存全占了。哪怕你的模型只需要2GB，它也可能把12GB全吞掉。为什么？因为它觉得"反正我以后可能要用，先占着再说"。

我个人习惯，从来不用默认配置。你想想看，一台机器上可能同时跑好几个实验，一个TF进程把显存全占了，其他实验怎么办？

所以，第一步就是手动控制显存分配。TF提供了两种方式：

按比例分配：设置每个GPU进程最多用多少比例的显存
按需分配（动态增长）：用多少占多少，不够了再申请

我建议你优先用动态增长策略。原因很简单——灵活。我在项目中遇到过好几次，按比例分配设了0.5，结果模型训练到一半显存不够了，直接崩掉。换成动态增长后，同样的模型反而能跑完。

4.2 动态增长策略实战

动态增长的配置代码其实就几行。但要注意，必须在创建任何TF操作之前设置。这个顺序我吃过亏，曾经把配置写在模型定义后面，结果完全不生效。

import tensorflow as tf

# 配置GPU显存动态增长
gpus = tf.config.experimental.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(f"配置失败: {e}")

这段代码跑完，TF就不会一开始就把显存全占了。它会随着你的模型运行，一点点申请显存。比如你的模型需要4GB，它就只占4GB左右，不会多占。

小提示： 动态增长虽然灵活，但有一个小代价——显存碎片化。如果你的模型反复创建和销毁大张量，显存可能会被切成很多小块，导致后续申请大块显存失败。我建议在训练脚本里只设置一次，不要频繁开关。

4.3 显存日志分析

光配置还不够，你得学会看日志。TF的显存日志，说白了就是告诉你"谁在吃显存、吃了多少、还剩多少"。

怎么开启显存日志？有两种方式：

环境变量法：设置 TF_CPP_VMODULE 来打印内存分配细节
回调函数法：在训练循环里手动打印显存使用情况

我个人更推荐第二种，因为可控。你可以在每个epoch结束后打印一次，不会刷屏。

import tensorflow as tf

class MemoryCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        # 获取当前显存使用情况
        gpu_memory = tf.config.experimental.get_memory_info('GPU:0')
        current = gpu_memory['current'] / (1024**3)  # 转成GB
        peak = gpu_memory['peak'] / (1024**3)
        print(f"Epoch {epoch+1}: 当前显存 {current:.2f}GB, 峰值 {peak:.2f}GB")

这段代码跑起来，你就能看到每个epoch的显存变化。如果发现峰值一直在涨，那大概率是存在显存泄漏。

注意： 显存泄漏在TF里其实挺常见的。尤其是你在循环里反复调用 tf.function 或者创建新的 tf.Variable 时，很容易把显存越吃越多。我曾经排查过一个案例，模型跑了200个epoch，显存从4GB涨到了11GB，最后直接OOM。后来发现是某个自定义层里每次前向都创建了一个临时变量，没有释放。

4.4 常见显存问题与排查思路

我整理了一张表，把TF显存相关的常见问题列了出来。你在排查时可以对照着看：

现象	可能原因	排查方法
一启动就占满显存	未开启动态增长	检查 `set_memory_growth` 是否生效
训练过程中显存持续增长	存在显存泄漏	用 `get_memory_info` 监控每个epoch的峰值
batch size调小后仍然OOM	模型本身占显存过大	检查是否有未释放的临时张量
多GPU训练时显存不均	数据分配不均衡	检查 `tf.distribute.MirroredStrategy` 的配置

嗯，这里要特别说一下显存泄漏的排查。我曾经遇到一个案例，模型跑着跑着显存就爆了。我一开始以为是batch size太大，调小了还是不行。后来用 get_memory_info 一查，发现每个epoch显存涨了200MB。顺着代码往下查，发现是一个自定义的损失函数里，每次计算都创建了一个 tf.constant，而且没有复用。改成全局变量后，问题就解决了。

4.5 知识体系总览

为了让你更直观地理解TF显存管理的全貌，我画了一张流程图：

这张图把TF显存管理的三个核心模块串起来了。你从左边开始看，先搞清楚配置策略，然后掌握动态增长的用法，最后学会用日志来排查问题。三个模块缺一不可。

  核心要点：
  永远不要用TF的默认显存配置，一定要手动设置
动态增长策略是首选，但要注意碎片化问题
显存日志是排查OOM的利器，每个epoch都看一眼
显存泄漏往往藏在自定义层或损失函数里

好了，TF显存管理这块就讲到这里。你回去可以试试把动态增长策略加到你的训练脚本里，再跑一遍看看显存变化。我保证，你会对"原来我的模型只吃这么点显存"感到惊讶。

📚 课程目录

4. TensorFlow显存管理：TF显存配置、动态增长策略、内存日志分析

4.1 TF的显存分配策略

4.2 动态增长策略实战

4.3 显存日志分析

4.4 常见显存问题与排查思路

4.5 知识体系总览