TensorFlow报OOM错误主因是默认预分配全部GPU显存，而非显存真被占满；应优先使用set_memory_growth或set_memory_limit限制显存，并检查tf.data管道、模型层及损失函数中的隐性显存泄漏。
为什么TensorFlow会报ResourceExhaustedError: OOM when allocating tensor
这不是显存真的被占满，而是TensorFlow默认“预分配全部可见GPU显存”，哪怕你只跑一个tf.constant([1.0])。NVIDIA驱动、其他进程（比如另一个Jupyter kernel）、甚至CUDA上下文初始化都会吃掉一部分显存，导致留给你的实际空间远小于标称值。更麻烦的是，TF 2.x的tf.function和自动图优化可能在后台悄悄放大中间张量——你没写循环，它自己缓存了10份梯度。
立即生效的显存限制方案（TF 2.x）
别碰allow_growth这种老方案，它不解决根本问题，反而让OOM更难复现。直接用set_memory_growth或set_memory_limit：
import tensorflow as tf
<p>gpus = tf.config.list_physical_devices('GPU')
if gpus:
try:</p><h1>方案A：按需增长（推荐调试时用）</h1><pre class="brush:php;toolbar:false;"><pre class="brush:php;toolbar:false;"> for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
# 方案B：硬性限制（训练时更稳，比如只给6GB）
# tf.config.experimental.set_memory_limit(gpus[0], 6144) # 单位MB
except RuntimeError as e:
print(e) # GPU已初始化后无法修改</code></pre>- 必须在
import tensorflow之后、任何<code>tf.*操作之前调用 -
set_memory_growth=True不是“省着用”，而是“每次需要多少就向系统申请多少”，避免初始霸占 -
set_memory_limit数值要留余量：显卡标称11GB ≠ 可用11GB，建议设为标称值的70%~80%
Batch size不是唯一变量：检查tf.data管道泄漏
OOM常发生在model.fit()启动后几轮才爆发，大概率是tf.data.Dataset里藏着隐患：
-
dataset.cache()把整个数据集加载进显存？错——它默认缓存在内存（RAM），但若你误加了.cache().prefetch(tf.data.AUTOTUNE)又没控制prefetch深度，缓冲区可能堆积大量未消费张量 -
map(..., num_parallel_calls=tf.data.AUTOTUNE)并行数过高，每个线程都持有一份预处理中间结果 - 自定义
map函数里用了tf.py_function，而Python端逻辑（如PIL读图+resize）没释放临时numpy数组
实操建议：
立即学习“Python免费学习笔记（深入）”；
# 安全写法：显式控制prefetch数量，禁用AUTOTUNE的不可控性 dataset = dataset.batch(32) dataset = dataset.prefetch(2) # 不用AUTOTUNE，固定为2层缓冲 dataset = dataset.map(preprocess_fn, num_parallel_calls=2) # 手动限并发
模型层与损失函数里的隐性显存炸弹
有些写法看着无害，实际每步都在复制张量：
-
tf.keras.layers.Reshape((-1,))本身不耗显存，但如果输入是[batch, H, W, C]且H*W*C极大，reshape只是视图变换，后续算子（如tf.matmul）会触发真实内存分配 - 自定义loss里写了
tf.print()或tf.debugging.assert_*，这些op在graph mode下仍会保留完整tensor用于输出/校验 - 用
tf.keras.Model封装时，把tf.Variable定义在call()里（而非__init__），每次call都新建变量 → 显存持续上涨
快速检测方法：运行前加这行，看每层输出shape是否合理：
model = build_model() model.build(input_shape=(None, 224, 224, 3)) model.summary() # 重点看"Param #"列是否异常大，以及output shape是否突然膨胀
显存问题最狡猾的地方在于：它不总在最大batch size下爆发，而是在某个特定数据样本（比如超长文本、超高分辨率图）进入pipeline时突然触发。所以别只测平均case，一定要用dataset.take(1).as_numpy_iterator()抽几条极端样本手动喂给model，观察nvidia-smi的显存跳变点。
