TensorFlow使用验证数据时的内存溢出问题排查

问题描述 #

我正在训练一个TensorFlow神经网络，因为训练数据太大超出电脑RAM，所以将其划分为子数据集进行序贯训练。不传入validation_data参数时，model.fit可以正常运行；但传入该参数后，程序在约第15轮epoch时抛出内存错误。

报错信息：

"C:\Users\PC\anaconda3\envs\CVISEnv\lib\site-packages\tensorflow\python\framework\constant_op.py", line 102, in convert_to_eager_tensor
return ops.EagerTensor(value, ctx.device_name, dtype)
numpy.core._exceptions._ArrayMemoryError: Unable to allocate 477. MiB for an array with shape (125000, 1000) and data type float32

疑惑：不是电脑无法读取验证数据，是不是每次调用model.fit都会复制验证数据？

问题原因与解答 #

你的疑惑是对的——如果直接传入numpy数组格式的validation_data，在循环调用model.fit进行序贯训练时，TensorFlow的Eager模式会在每次调用时将验证数据重新转换为EagerTensor，加上训练过程中模型参数优化、中间张量的内存占用累积，到第15轮时内存就被耗尽了。而且numpy数组本身是一次性加载到内存的，多次处理会导致冗余的内存占用叠加。

解决方案 #

1. 将验证数据转为tf.data.Dataset格式（推荐） #

tf.data.Dataset是按需加载数据的，不会一次性把所有验证数据都驻留在内存，而且每次调用model.fit时不会重复复制整个数据集，内存占用更稳定。示例代码：

import tensorflow as tf

# 假设你的验证数据是numpy数组val_x和val_y
val_dataset = tf.data.Dataset.from_tensor_slices((val_x, val_y))
val_dataset = val_dataset.batch(32)  # 根据你的需求设置batch大小

之后在model.fit中传入validation_data=val_dataset即可。

2. 手动清理内存（备选） #

如果必须使用numpy数组，每次调用model.fit后可以手动清理内存，减少累积：

import gc
import tensorflow as tf

# 训练完一个子数据集后
model.fit(...)
tf.keras.backend.clear_session()
gc.collect()

注意：clear_session会重置Keras会话，如果你需要保留模型状态，这个方法要谨慎使用。

3. 优化验证数据加载逻辑 #

确保验证数据只在循环外加载一次，不要在每次序贯训练的循环内重复读取验证数据，避免不必要的内存复制。

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