tensorflow 默认预分配全部gpu显存，需在import后立即调用set_memory_growth(true)，且必须在任何tf操作前生效；cache()应置于轻量预处理后、batch前，并配合prefetch(autotune)避免显存爆炸。

这不是模型或代码写错了，而是 TensorFlow 默认把所有可见 GPU 显存预分配并锁死——哪怕你只执行 tf.constant(1)，nvidia-smi 也会显示显存“占满”，但 GPU-Util 始终徘徊在个位数。

set_memory_growth 必须在 import 后立刻调用

这个配置不是“开关”，而是一次性初始化指令：它必须在任何 TensorFlow 操作（包括 tf.data.Dataset 构建、tf.keras.Model() 实例化、甚至 tf.config.list_physical_devices('GPU') 之后的第一次张量创建）之前生效。一旦 GPU 设备被内部初始化，再调用就完全无效。

• os.environ['CUDA_VISIBLE_DEVICES'] 必须在这之前设置，否则 tf.config.experimental.list_physical_devices('GPU') 可能返回空列表
• 多卡时不能只设 gpus[0]，要遍历所有 gpus 并对每个调用 set_memory_growth(True)
• TF 2.12+ 中混用 tf.compat.v1.ConfigProto 和 Session 是无效的，不会报错但也不起作用

cache() 放错位置 = 显存爆炸开关

.cache() 不是“缓存一下”，而是把上游整个 pipeline 的输出一次性加载进显存（或内存），且 epoch 结束也不释放。它不看 batch 多小、模型多轻，只认“上游有没有做完预处理”。

python全能编程助手

SkillSub Pro - Python 题解与代码注释双功能技能功能概述SkillSub Pro - Python 题解与代码注释双功能技能是一项面向实际任务的技能，主要用于SkillSub Pro 是一个 Python 题解生成与代码注释的 双功能合体技能 ，专为学生、算法学习者和开发者设计；✅ 一个技能，两种用途 ：；核心要点📝 题解模式 ：输入题目/题号，自动生成完整 Python 题解（含详细注释、解题思路、复杂度分析）；💬 注释模式 ：输入 Python 代码，自动添加详细中。它将相关步骤、

下载

• 错误写法：.map(preprocess_fn).cache().shuffle().batch() → 预处理前就缓存原始大图，显存瞬间飙高
• 正确顺序：.map(light_preprocess).cache().shuffle().batch().prefetch(tf.data.AUTOTUNE)，其中 light_preprocess 仅做降分辨率、裁剪、归一化等轻量操作
• 实在要缓存大图，改用磁盘缓存：.cache('/path/to/disk/cache')，首次慢但彻底避开显存

数据管道卡在 CPU → GPU 空转

显存被占满但 GPU-Util 接近 0，大概率是 CPU 拿不出数据给 GPU 算——nvidia-smi 看到的是“已分配显存”，不是“活跃计算”。此时 GPU 在等，CPU 在忙磁盘 I/O 或慢速 Python 预处理。

• 检查 htop 或任务管理器：如果 CPU 使用率持续 90%+，说明瓶颈在数据加载
• 避免 from_generator + 复杂 Python 函数；优先用 tf.io.decode_jpeg 等原生算子
• 用 tf.data.TFRecordDataset 替代原始文件读取，绕过 CPU 解码瓶颈
• .prefetch(tf.data.AUTOTUNE) 必须有，且应放在 .batch() 之后，否则预取失效

真正难调试的点在于：显存“占满”和“用满”是两回事。nvidia-smi 的 Memory-Usage 数值本身不反映实际张量活跃度，很多是框架预留或未释放的缓冲区。先确认是不是 set_memory_growth 没生效、cache() 放太前、或数据管道堵在 CPU，再动模型结构或 batch size。

Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！