第3章：TensorFlow显存管理：GPU配置选项、显存增长策略

说到TensorFlow的显存管理，我刚开始接触时也踩过不少坑。你想想看，模型明明不大，一跑起来却提示OOM（Out of Memory），这谁受得了？其实问题往往出在默认配置上——TensorFlow默认会“霸占”整张GPU的显存，哪怕你只需要一点点。

这一章，我们就来聊聊怎么管好这块“肥肉”。

3.1 GPU配置选项：从“全占”到“按需”

TensorFlow 2.x 提供了 tf.config 这个配置接口。我个人习惯在代码开头就把GPU策略定好，免得跑了一半才发现显存不够。

先看一个最基础的配置：

import tensorflow as tf

# 列出所有物理GPU设备
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    print(f"检测到 {len(gpus)} 个GPU")
    for gpu in gpus:
        print(f"  - {gpu.name}")
else:
    print("没有检测到GPU，请检查环境")

这段代码会告诉你当前机器上有几张卡、叫什么名字。我在项目中遇到过一台服务器插了4张RTX 3090，结果代码默认只用了第0张，其他三张闲着。嗯，这时候就需要手动指定了。

3.1.1 指定可见GPU

有时候我们只想用某几张卡，可以用 tf.config.set_visible_devices：

# 只使用第0号和第1号GPU
tf.config.set_visible_devices(gpus[0:2], 'GPU')

为什么要这么做？说白了，多卡训练时，如果代码没写好，反而会因为通信开销拖慢速度。我建议先单卡调通，再上多卡。

3.2 显存增长策略：set_memory_growth

这是本章的重头戏。默认情况下，TensorFlow会一次性占满所有显存。比如你有一张12GB的卡，它直接吃掉11.9GB，其他程序根本没法用。

解决方案就是 显存增长（Memory Growth）。

核心概念： 让TensorFlow按需分配显存，用多少占多少，而不是一开始就全占。

看代码：

gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        # 为每个GPU开启显存增长
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        print("显存增长模式已开启")
    except RuntimeError as e:
        print(f"设置失败: {e}")

这段代码我几乎每个项目都会写。你想想看，如果服务器上同时跑两个实验，一个占满显存，另一个就只能干瞪眼。开启增长后，两个实验可以共存，各取所需。

小提示： set_memory_growth 必须在任何TensorFlow操作之前调用。一旦创建了张量或模型，这个设置就失效了。我习惯把它放在 import tensorflow 之后的第一行。

3.2.1 显存增长的底层逻辑

为什么会这样？TensorFlow底层用的是CUDA的显存分配器。默认策略是“先占满再说”，而增长策略是“用一点，申请一点”。

具体来说：

默认模式：启动时调用 cudaMalloc 分配几乎所有显存
增长模式：按需调用 cudaMalloc，每次分配一小块

我在项目中遇到过一个问题：开启增长后，训练速度反而变慢了。排查后发现，是因为频繁的显存分配导致开销增加。解决办法是配合 tf.config.experimental.set_memory_limit 设置一个上限。

3.3 显存限制：给GPU“划地盘”

除了增长模式，我们还可以直接限制最大显存用量。比如一张12GB的卡，我只想用4GB：

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

这里 memory_limit 的单位是MB。4096就是4GB。这样做的好处是：

给其他程序留出空间（比如数据预处理、可视化）
防止单个实验吃掉所有资源
方便在同一张卡上跑多个小模型

警告： 设置显存限制后，如果模型实际需要超过限制，会直接报OOM错误。我曾经在调参时把限制设得太小，结果模型死活跑不起来，排查了半天才发现是显存限制的问题。

3.4 实战：完整配置模板

下面是我个人常用的配置模板，你可以直接拿去用：

import tensorflow as tf

# 1. 检测GPU
gpus = tf.config.list_physical_devices('GPU')
if not gpus:
    raise RuntimeError("没有检测到GPU")

# 2. 设置可见GPU（这里用第0张）
tf.config.set_visible_devices(gpus[0], 'GPU')

# 3. 开启显存增长
tf.config.experimental.set_memory_growth(gpus[0], True)

# 4. 可选：设置显存上限（比如6GB）
# tf.config.set_logical_device_configuration(
#     gpus[0],
#     [tf.config.LogicalDeviceConfiguration(memory_limit=6144)]
# )

print("GPU配置完成")
print(f"使用GPU: {gpus[0].name}")

这段代码我几乎每个项目都复用。你想想看，有了这个模板，不管换到哪台机器，显存管理都不会出大问题。

3.5 知识体系总览

为了让你更直观地理解本章的知识结构，我画了一张图：

这张图把本章的三个核心知识点串起来了。你想想看，从配置选项到增长策略，再到显存限制，其实是一条线：先知道有哪些卡，再决定怎么用，最后设个上限保底。

3.6 避坑指南

最后，分享几个我踩过的坑：

曾经我在Jupyter Notebook里反复运行 set_memory_growth，结果第二次运行时报错。原因是TensorFlow在第一次运行后已经初始化了GPU上下文，不能再改配置。解决办法是重启kernel，或者把配置写在最前面。
曾经我在Docker容器里忘了映射GPU设备，结果代码怎么都检测不到GPU。嗯，记得加 --gpus all 参数。
曾经我同时开了显存增长和显存限制，结果发现限制优先级更高。说白了，增长是“软策略”，限制是“硬约束”。

好了，这一章的内容就到这里。显存管理看似小事，但搞不好真的会让人抓狂。希望这些经验能帮你少走弯路。

📚 课程目录

第3章：TensorFlow显存管理：GPU配置选项、显存增长策略

3.1 GPU配置选项：从“全占”到“按需”

3.1.1 指定可见GPU

3.2 显存增长策略：set_memory_growth

3.2.1 显存增长的底层逻辑

3.3 显存限制：给GPU“划地盘”

3.4 实战：完整配置模板

3.5 知识体系总览

3.6 避坑指南