在Kaggle中降低机器学习训练内存占用的方法求助

问题描述 #

我在Kaggle训练机器学习模型，数据集是一个形状为(34792, 369, 103, 4)的NumPy数组，包含34792张369×103像素的CMYK格式图像。目前能正常导入数据集并划分训练集、测试集，但在K折交叉验证调用fit方法时，内存占用过高导致笔记本崩溃。已经尝试将所有数据转为uint8类型、删除无用数组并调用垃圾回收机制，问题仍未解决。

现有代码 #

kfold = KFold(n_splits=10, shuffle=True, random_state=42)

for train, val in kfold.split(X_train):
    model = keras.Sequential()
    model.add(keras.layers.Conv2D(64, (3,3), activation='relu', input_shape=X_shape))
    model.add(keras.layers.MaxPooling2D((2,2)))
    model.add(keras.layers.Conv2D(64, (3,3), activation='relu'))
    model.add(keras.layers.MaxPooling2D((2,2)))
    model.add(keras.layers.Conv2D(64, (3,3), activation='relu'))
    model.add(keras.layers.MaxPooling2D((2,2)))
    model.add(keras.layers.Conv2D(64, (3,3), activation='relu'))
    model.add(keras.layers.MaxPooling2D((2,2)))
    model.add(keras.layers.Flatten())
    model.add(keras.layers.Dense(128, activation='relu'))
    model.add(keras.layers.Dense(1, activation='sigmoid'))

optimizer = keras.optimizers.Adam(learning_rate = 0.001)
model.compile(optimizer="adam",loss='binary_crossentropy', metrics=['accuracy'])


early_stop = keras.callbacks.EarlyStopping(monitor='val_loss', restore_best_weights= True, patience=5, verbose=1)
callback = [early_stop]


history = model.fit(X_train[train], y_train[train], epochs = 50, validation_data=(X_train[val], y_train[val]), batch_size=64, callbacks=callback)


scores = model.evaluate(X_train[val], y_train[val], verbose=0)
print(scores)

降低内存占用的可行方法 #

• 图像格式压缩：CMYK转RGB/灰度图
  CMYK是4通道，转成RGB可减少1/4内存；若任务允许，转单通道灰度图直接将内存占用降至原25%，这是最直接的压缩手段。示例代码：

  # 转灰度图（取通道平均值）
  X_train_gray = np.mean(X_train, axis=-1, keepdims=True)
  
  # 用PIL转RGB（精准色彩转换）
  from PIL import Image
  def cmyk_to_rgb(cmyk_img):
      img = Image.fromarray(cmyk_img, mode='CMYK')
      return np.array(img.convert('RGB'))
  

• 用生成器/TF Dataset分批加载数据
  避免一次性加载全量数据，改用分批加载方式，仅在训练时加载当前batch的样本：

  import tensorflow as tf
  
  # 用tf.data.Dataset实现高效分批
  def create_dataset(X, y, batch_size):
      dataset = tf.data.Dataset.from_tensor_slices((X, y))
      dataset = dataset.shuffle(len(X)).batch(batch_size).prefetch(tf.data.AUTOTUNE)
      return dataset
  
  # K折循环内调用
  train_dataset = create_dataset(X_train[train], y_train[train], 64)
  val_dataset = create_dataset(X_train[val], y_train[val], 64)
  
  history = model.fit(train_dataset, 
                      epochs=50, 
                      validation_data=val_dataset,
                      callbacks=callback)
  

• 减小Batch Size
  当前batch_size=64，可尝试降至32或16，直接减少单轮迭代的内存占用，牺牲少量训练速度换取内存空间。

• 模型轻量化改造

  • 减少卷积过滤器数量：将Conv2D(64, ...)改为Conv2D(32, ...)，降低模型参数和中间特征图内存消耗；
  • 替换为深度可分离卷积：用SeparableConv2D替代普通Conv2D，参数和计算量仅为原1/10左右；
  • 精简模型结构：尝试减少一层卷积层，观察性能变化的同时降低内存压力。

• K折循环内强制清理内存
  每次折训练结束后，显式删除模型并触发垃圾回收，避免内存累积：

  import gc
  
  for train, val in kfold.split(X_train):
      # 创建、训练模型逻辑...
      
      # 训练后清理
      del model
      gc.collect()
  

• 缩小图像尺寸
  若任务对分辨率要求不高，将图像缩至原尺寸的1/2或1/3，比如369×103→184×51，单张图像内存占用降至原1/4。示例：

  def resize_img(img, target_size=(184,51)):
      img_pil = Image.fromarray(img, mode='CMYK')
      return np.array(img_pil.resize(target_size))
  

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