如何修复Python中TensorFlow显存溢出导致的OOM错误？-Python教程-PHP中文网

tf.keras.model.fit()报oom错误主因是tensorflow默认预分配全部gpu显存，叠加batch过大、图像尺寸高、中间层未释放及tf.data.cache()缓存全量数据所致；应立即调用tf.config.experimental.set_memory_growth()启用动态显存分配，并同步调小batch size、降低输入分辨率、优化tf.data流水线与gradienttape使用。

如何修复python中tensorflow显存溢出导致的oom错误？

为什么`tf.keras.Model.fit()`突然报`ResourceExhaustedError: OOM when allocating tensor`

这不是模型真的需要那么多显存，而是TensorFlow默认把整块GPU显存预分配满（尤其在1.x和2.x早期版本），再叠加batch过大、输入图像尺寸太高、模型中间层输出没及时释放，就直接触发OOM。常见于多卡环境没设CUDA_VISIBLE_DEVICES，或用了tf.data.Dataset.cache()把整个数据集缓存进显存。

用`tf.config.experimental.set_memory_growth()`动态分配显存

这是最常用也最有效的缓解手段，告诉TensorFlow“按需申请，别一口吞完”。注意它必须在导入tensorflow后**立即调用**，且只对当前GPU生效：

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)
    except RuntimeError as e:
        print(e)  # 初始化后不能再修改

不支持Windows上的CUDA 11.2+与TF 2.8+组合（会报Failed to enable memory growth）
若用tf.distribute.MirroredStrategy，得在strategy创建前调用
设了这个之后，nvidia-smi看到的显存占用会“缓慢上涨”，而非启动即占满

控制batch size和输入尺寸比调显存参数更直接

显存消耗≈batch_size × height × width × channels × model_depth × sizeof(float32)。很多用户花一小时调set_memory_limit()，不如把batch_size=32改成16，或把512×512图缩到384×384——立刻见效。

A Python CLI skill for Cutout.Pro visual APIs — background removal, face cutout, and photo enhancement. Supports file upload & image URL input.

调用 Cutout.Pro 视觉处理 API 进行背景移除、人像抠图和照片增强，支持文件上传与图片 URL 输入。

下载

用tf.data.Dataset.batch(batch_size, drop_remainder=True)避免最后一个小batch导致形状不一致
训练时用tf.image.random_crop()代替全图加载，能显著降低单步峰值显存
验证阶段可临时把batch_size设为训练时的2–4倍（因无梯度计算），但要注意model.evaluate()内部仍会按batch分片，不是真省显存

`tf.function`里嵌套`tf.GradientTape`容易隐式累积显存

自定义训练循环中，如果在@tf.function装饰的函数里反复新建tf.GradientTape(persistent=False)，TensorFlow可能无法及时回收中间张量。尤其当tape包裹了model(x)又调了多次tape.gradient()，会把所有前向中间结果都留着。