为何运行1GB鸟类CNN数据集的train_test_split时内核总是崩溃

你当前的实现是将所有图像一次性加载为numpy数组后再做拆分，原始1GB的图像数据集转成未压缩的numpy数组后，实际内存占用会达到原始文件大小的3~5倍，再加上train_test_split执行时会生成数组副本，超出可用内存就会触发内核崩溃。你设置shuffle=False没有效果，因为问题本质不是shuffle的计算开销，是内存不足。

• 优先改用流式加载API，避免全量数据入内存
  你使用TensorFlow框架的话，直接调用tf.keras.utils.image_dataset_from_directory接口，直接从文件夹按分类加载数据，自动完成拆分、预处理，全程不需要把所有图像加载到内存，内存占用可以降低90%以上，示例用法如下：

  import tensorflow as tf
  img_size = (100, 100)
  batch_size = 32
  # 直接加载训练集并按8:2拆分训练和验证集
  train_ds = tf.keras.utils.image_dataset_from_directory(
    'birds/train',
    validation_split=0.2,
    subset="training",
    seed=42,
    image_size=img_size,
    batch_size=batch_size
  )
  val_ds = tf.keras.utils.image_dataset_from_directory(
    'birds/train',
    validation_split=0.2,
    subset="validation",
    seed=42,
    image_size=img_size,
    batch_size=batch_size
  )
  # 归一化
  normalization_layer = tf.keras.layers.Rescaling(1./255)
  train_ds = train_ds.map(lambda x, y: (normalization_layer(x), y))
  val_ds = val_ds.map(lambda x, y: (normalization_layer(x), y))
  

• 优化现有代码的内存占用
  如果你要保留当前手动读数据的逻辑，按以下步骤调整即可：

  1. 先拆分索引再取数，不要直接传入全量图像数组：先对样本索引做拆分，再用索引从img_data中取对应数据，避免train_test_split内部复制整份图像数组
  2. 拆分完成后再做归一化：uint8类型的图像数组内存占用仅为float32类型的1/4，拆分前不要做除以255的类型转换
  3. 你已经手动做了全局shuffle，可以直接用下标切分替代train_test_split，完全避开sklearn接口的额外内存开销：

  # 直接按比例切分，不需要调用train_test_split
  split_point = int(len(img_data) * 0.8)
  X_train, X_test = img_data[:split_point], img_data[split_point:]
  y_train, y_test = img_labels[:split_point], img_labels[split_point:]
  # 切分后再做归一化
  X_train = X_train / 255.0
  X_test = X_test / 255.0
  

• 修复现有代码的显性bug
  你当前代码缺少import numpy as np语句，且拆分时只生成了X_train、X_test，后续直接使用未定义的X_val变量，会直接触发运行报错。

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