TensorFlow为何将全部数据加载至系统内存？OOM错误求助

问题背景 #

持续遭遇系统内存（非GPU内存）OOM错误，代码基于之前的图像分类器修改，仅做少量调整：

• 原始图像尺寸更大，但已提前调整为224x224；
• 数据集规模翻倍，且已移除cache和shuffle操作，但仍未按batch加载数据，在第一个epoch开始前崩溃。

核心原因分析 #

1. from_tensor_slices的隐式内存加载：如果img_paths或oh_input是大型numpy数组/普通列表，tf.data.Dataset.from_tensor_slices会将整个数据集一次性加载到系统内存中，这是导致OOM的首要原因。
2. 全量数据集的shuffle缓冲区：在ds_split函数中，shuffle_size传入了len(img_paths)，即使用整个数据集作为shuffle缓冲区，会强制把所有数据加载到内存完成打乱操作。
3. map操作的串行执行：默认的map是单线程处理，可能导致数据预处理过程中内存堆积。

具体解决步骤 #

1. 替换from_tensor_slices为延迟加载方式 #

避免一次性加载所有路径和标签到内存，改用tf.convert_to_tensor或生成器方式：

# 用tf.convert_to_tensor避免numpy数组的内存拷贝
img_paths_tensor = tf.convert_to_tensor(img_paths, dtype=tf.string)
oh_input_tensor = tf.convert_to_tensor(oh_input, dtype=tf.float32)
ds_oh = tf.data.Dataset.from_tensor_slices((img_paths_tensor, oh_input_tensor))

如果oh_input是极大数组，建议用生成器进一步降低内存占用：

def data_generator():
    for path, label in zip(img_paths, oh_input):
        yield path, label

ds_oh = tf.data.Dataset.from_generator(
    data_generator,
    output_signature=(
        tf.TensorSpec(shape=(), dtype=tf.string),
        tf.TensorSpec(shape=(29,), dtype=tf.float32)
    )
)

2. 缩小shuffle缓冲区大小 #

不要用整个数据集大小作为shuffle_size，改为合理数值（如1000-5000，根据内存情况调整）：

# 修改ds_split调用，shuffle_size设为1000
train_ds, val_ds = ds_split(ds_oh, len(img_paths), 1000, train_split=0.8, val_split=0.2, shuffle=True)

同步调整ds_split函数内的shuffle逻辑：

def ds_split(ds, ds_size, shuffle_size, train_split=0.8, val_split=0.2, shuffle=True):
    assert (train_split + val_split) == 1
    
    if shuffle:
        # 使用指定的shuffle_size，而非全量数据集
        ds = ds.shuffle(shuffle_size, seed=99)
    
    train_size = int(train_split * ds_size)
    val_size = int(val_split * ds_size)
    
    train_ds = ds.take(train_size)    
    val_ds = ds.skip(train_size).take(val_size)
    
    return train_ds, val_ds

3. 开启map的并行处理 #

在map操作中添加并行参数，提升预处理效率的同时避免内存堆积：

ds_oh = ds_oh.map(read_and_decode, num_parallel_calls=tf.data.AUTOTUNE, deterministic=False)

4. 优化数据类型减少内存占用 #

在read_and_decode中延迟数据类型转换，减少内存占用：

def read_and_decode(filename, label):
    img = tf.io.read_file(filename)
    img = tf.io.decode_jpeg(img, channels=3)
    img = tf.image.resize_with_pad(img, 224, 224, method=tf.image.ResizeMethod.BILINEAR)
    # 延迟转换为float32，减少内存占用
    img = tf.cast(img, tf.float32)
    img = preprocess_input(img)
    return img, label

5. 验证数据集迭代逻辑 #

测试只取一个batch，确认内存变化，定位问题阶段：

# 测试只取一个batch，查看内存变化
for batch in train_ds.take(1):
    print(batch[0].shape, batch[1].shape)

如果这一步就OOM，说明问题出在数据集构建阶段，而非模型训练。

内容的提问来源于stack exchange，提问作者John G.