tf.keras训练模型遇OOM错误：为何按全量数据分配张量？

为什么会出现这个OOM错误？ #

你遇到的核心问题是输入数据的类型导致TensorFlow没有按你设置的batch_size进行分批处理，而是尝试一次性处理整个900000样本的数据集。

具体来说，你的X_train是一个TensorFlow Variable对象，当直接将Variable传入model.fit()时，TensorFlow在构建计算图阶段会默认把整个张量视为完整输入，而非按批次拆分。再加上你启用了use_multiprocessing=True，多进程机制会进一步触发对整个数据集的内存预分配——这就是为什么你看到系统尝试分配[900000,28,28,20]这样的超大张量，而不是预期的[32,28,28,20]批次张量。

解决方法 #

这里有几个可行的修复方案，按推荐程度排序：

1. 使用tf.data.Dataset封装输入数据（最推荐）

将你的训练数据转换为tf.data.Dataset，这是TensorFlow官方推荐的输入数据管道方式，能确保严格按批次加载数据，避免一次性占用大量GPU内存：

# 将Variable转换为Dataset，并按batch_size分批
train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train)).batch(BATCH_SIZE)

# 使用Dataset训练，无需指定batch_size和steps_per_epoch（可选指定）
model.fit(train_dataset, epochs=5, use_multiprocessing=True)

如果你仍然想手动指定steps_per_epoch，可以保持原来的计算方式，但建议优先让Dataset自动处理批次。

2. 将Variable转换为普通NumPy数组

如果你的CPU内存足够容纳整个数据集，可以把X_train和y_train转换为NumPy数组，这样model.fit()会自动按批次拆分：

X_train_np = X_train.numpy()
y_train_np = y_train.numpy()

model.fit(X_train_np, y_train_np, batch_size=BATCH_SIZE, steps_per_epoch=int(X_train.shape[0])//BATCH_SIZE, epochs=5, use_multiprocessing=True)

这种方式简单直接，但如果数据集过大（比如你的90万样本），可能会占用较多CPU内存。

3. 禁用use_multiprocessing=True（临时应急方案）

多进程模式在处理Tensor类型输入时容易出现内存预分配问题，暂时禁用它可以避免一次性加载整个数据集：

model.fit(X_train, y_train, batch_size=BATCH_SIZE, steps_per_epoch=int(X_train.shape[0])//BATCH_SIZE, epochs=5)

但这会损失多进程带来的训练加速效果，所以仅作为临时解决方法。

额外建议 #

训练前可以通过以下代码开启GPU内存动态增长，避免TensorFlow一次性占满所有GPU内存：

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)

内容的提问来源于stack exchange，提问作者Rahul