如何优化训练集图像特征向量提取？解决Kaggle内存不足问题

原代码的核心问题是一次性加载所有15k张图像到内存，导致内存占用过载。以下是针对性的优化方案：

优化思路：分批处理 #

通过分批次读取、预处理、提取特征，避免一次性占用大量内存，最后再合并所有批次的结果。

优化后的代码 #

def save_trained_file_feature_vectors(feat_extractor, save_path, batch_size=32):
    all_features = []
    
    # 按批次遍历图像文件
    for i in range(0, len(train_files), batch_size):
        batch_files = train_files[i:i+batch_size]
        batch_images = []
        
        # 加载当前批次的图像
        for f in batch_files:
            original = load_img(f, target_size=(imgs_model_width, imgs_model_height))
            numpy_image = img_to_array(original)
            batch_images.append(numpy_image)
        
        # 预处理并提取特征
        batch_array = np.array(batch_images)
        processed_batch = preprocess_input(batch_array)
        batch_features = feat_extractor.predict(processed_batch, verbose=0)
        
        all_features.append(batch_features)
        
        # 手动释放当前批次的内存
        del batch_images, batch_array, processed_batch, batch_features
    
    # 合并所有批次特征并保存
    imgs_features = np.vstack(all_features)
    np.save(save_path, imgs_features)
    return imgs_features

额外优化建议 #

• 调整批次大小：根据Kaggle的可用内存（通常16GB左右），可以尝试把batch_size调到64或128，找到不触发内存溢出的最大值。
• 用Keras生成器自动管理内存：如果图像按文件夹分类，直接用ImageDataGenerator生成批次数据，无需手动处理：

  from tensorflow.keras.preprocessing.image import ImageDataGenerator
  
  def save_trained_file_feature_vectors(feat_extractor, save_path, img_dir, batch_size=32):
      datagen = ImageDataGenerator(preprocessing_function=preprocess_input)
      generator = datagen.flow_from_directory(
          img_dir,
          target_size=(imgs_model_width, imgs_model_height),
          batch_size=batch_size,
          class_mode=None,  # 仅加载图像，不需要标签
          shuffle=False  # 保持特征与文件顺序一致
      )
      
      imgs_features = feat_extractor.predict(generator, verbose=1)
      np.save(save_path, imgs_features)
      return imgs_features
  

• 清理内存碎片：在批次处理后加入gc.collect()强制回收内存，减少内存残留：

  import gc
  gc.collect()
  

• 去掉不必要的复制：原代码中preprocess_input(images.copy())可改为preprocess_input(images)，避免额外的内存占用。

关于Dask无效的说明 #

Dask适合超大规模数据集，但如果仍采用一次性加载数据的方式，自然起不到效果。若要使用Dask，需分块读取图像并适配Keras模型，但实现复杂度远高于手动分批或生成器方案，没必要舍近求远。

内容的提问来源于stack exchange，提问作者Khoa Hữu Bảo Huỳnh