TensorFlow模型训练：ipynb运行正常.py文件执行崩溃如何解决

你调小batch size无效是因为本次崩溃的核心原因是CPU系统内存不足，不是GPU显存不足：

1. 报错日志明确显示CPU内存分配失败、且在填充Shuffle Buffer时进程被系统杀死，Shuffle Buffer占用的是CPU内存，和batch size关联度极低。
2. GPU显存仅占用2.5GB说明数据没有正常被加载到GPU进行计算，全部堆积在CPU内存侧导致占用飙升。
3. Jupyter Notebook和.py文件运行表现不同的常见原因是：Notebook分段执行时会自动释放部分临时变量内存，且部分Notebook环境默认开启TensorFlow内存优化配置，而.py文件全量执行时临时变量没有及时释放、配置未同步导致内存占用过高。

1. 优化数据集Shuffle配置 #

• 不要将shuffle(buffer_size)的buffer_size参数设置为全数据集大小，调整为单batch size的10~100倍即可（比如batch size为32时，buffer_size设为1024就可以保证足够的打乱效果），过大的buffer会直接占满CPU内存。
• 如果必须要大粒度的打乱效果，可以增加参数reshuffle_each_iteration=False，避免每轮epoch都重新加载全量buffer占用内存。

2. 同步TensorFlow内存配置 #

在.py文件导入TensorFlow后立刻添加显存增长配置，避免TensorFlow默认占满全部显存同时将数据堆积到CPU：

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

如果数据流水线用了prefetch(tf.data.AUTOTUNE)，先改成固定小值比如prefetch(2)，部分环境下AUTOTUNE会过度预加载数据占满CPU内存。

3. 优化数据集加载逻辑 #

• 检查.py代码中是否存在将全量训练数据直接加载为全局numpy数组/列表的逻辑，改为用tf.data流水线按需加载，每次只读取当前需要的batch数据。
• 训练过程中不要存储每一步的loss、准确率等全部指标到内存，只保留滑动平均指标或者定期写入日志文件即可，避免指标缓存占用过多内存。

4. 对齐运行环境 #

• 确认运行.py文件的Python环境和Jupyter使用的是同一个环境，避免TensorFlow版本差异导致的内存策略不同。
• 直接用命令行执行.py文件，不要用IDE的调试模式运行，调试模式会额外缓存大量运行数据，提升内存占用。

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