为何运行这段Python数据拆分代码后Google Colab内存占用超半数？

问题根源 #

你用train_test_split拆分数据时，默认会生成原始数据的副本，加上拆分了两次，内存里同时存着原始的ls_image、ls_labels，还有拆分后的feature、feature_test、feature_train、feature_val等多份数据——等于把几万条数据重复存了好几遍，自然会占掉大量内存。另外你代码里第二次拆分用的label是笔误，应该是label_。

实用解决方案 #

• 及时删除无用变量并回收内存
  拆分完成后立刻删掉不再需要的原始数据和中间变量，让Colab释放内存：

  import gc
  
  seed_constant = 123
  # 第一次拆分
  feature, feature_test, label_, label_test= train_test_split(ls_image, ls_labels, test_size=0.2, train_size=0.8)
  # 第二次拆分（修正笔误）
  feature_train, feature_val, class_train, class_val = train_test_split(feature, label_, test_size=0.25, train_size=0.75)
  
  # 删除原始数据和中间变量
  del ls_image, ls_labels, feature, label_
  gc.collect()  # 强制触发内存回收
  

• 改用分批加载数据的方式
  别一次性把所有图片读进内存存成ls_image，直接从文件夹分批读取，内存里只保留当前批次的图片。比如用TensorFlow的工具：

  import tensorflow as tf
  
  # 加载训练+验证集
  train_val_ds = tf.keras.utils.image_dataset_from_directory(
      '你的数据文件夹路径',
      validation_split=0.2,
      subset="training",
      seed=seed_constant,
      image_size=(224, 224),  # 改成你模型需要的尺寸
      batch_size=32)
  # 加载测试集
  test_ds = tf.keras.utils.image_dataset_from_directory(
      '你的测试数据文件夹路径',
      image_size=(224, 224),
      batch_size=32)
  # 从训练+验证集里拆分出验证集
  val_batches = tf.data.experimental.cardinality(train_val_ds)
  val_ds = train_val_ds.take(val_batches // 4)
  train_ds = train_val_ds.skip(val_batches // 4)
  

• 缩小图片分辨率
  如果你的图片尺寸太大（比如1000x1000以上），哪怕只有几千张，总内存也会飙升。先把所有图片resize到模型需要的输入尺寸（比如224x224），单张图片的内存占用会大幅降低。

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