使用TensorFlow存储4D NumPy矩阵时触发段错误的原因

问题基本信息 #

• 业务流程：生成4维NumPy矩阵后，将其导出为TensorFlow可加载的持久化数据集
• 数据规格：矩阵形状为[14540,512,128,4]，元素类型为np.float32
• 内存前提：确认数据完全可容纳在80GB DDR4内存中，不关注运行时内存占用
• 实现逻辑：

matrix = np.load(output_path+".npy")
if seed is None:
    seed = np.random.randint(1e6)

print("unshuffled")
if output_path is None:
    return tf.data.Dataset.from_tensor_slices(matrix)
else:
    tf.data.experimental.save(
        tf.data.Dataset.from_tensor_slices(matrix),
        output_path+"_unshuffled",
        compression="GZIP"
    )

• 功能目标：基于上述流程分别存储未打乱、打乱两个版本的数据集，seed参数仅用于后续的数据集打乱操作
• 异常表现：程序运行时抛出段错误，内存占用攀升至约22GB时即触发错误，远低于80GB的总物理内存上限

问题核心诱因 #

1. tf.data.Dataset.from_tensor_slices 的全量内存拷贝机制
   该接口不会直接复用NumPy数组的内存块，会将传入的数组完整拷贝到TensorFlow的张量内存域。按照数据规模计算，单份目标矩阵本身占用约14.2GB内存：np.load完成数组加载后已经占掉14.2GB内存，调用from_tensor_slices时会再申请同等大小的内存完成拷贝，两者叠加的内存规模已经接近28GB。内存申请是逐块分配的，当攀升到22GB区间时，如果出现内存碎片化、地址空间不足的问题，就会直接触发段错误，而非抛出可捕获的Python层级异常。
2. GZIP压缩的临时缓冲内存开销
   开启compression="GZIP"参数时，TensorFlow的持久化写入流程会为每个待写入的分片申请临时压缩缓冲区，且不会提前对这些临时内存的总峰值做校验。并行写入时多个分片的压缩缓冲区叠加，会进一步推高内存峰值，这类底层C++层的内存申请失败，往往不会向上传递Python异常，会直接触发进程段错误。
3. 环境层面的单进程内存寻址上限限制
   哪怕物理内存有80GB，如果运行环境存在部分32位依赖组件（比如旧版NumPy、TensorFlow的编译链接库未完成全64位适配），单进程的虚拟内存寻址上限往往被限制在20-30GB区间。进程内存申请超过该阈值时，会因无权限访问对应地址空间直接触发段错误，这也和你观察到的“22GB左右触发错误”的特征高度匹配。

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