大数据集交叉验证的内存优化方案咨询

我有一个约200GB的特征数据集（存储为Hugging Face生成的.arrow文件），需要基于它做随机森林模型的交叉验证。现在就算用600GB内存，代码还是会崩溃。我觉得问题出在把Hugging Face Dataset转成Pandas DataFrame再转Numpy数组的流程（代码前5行），但没找到替代方案。另外也想了解其他提升代码内存效率的方法。

原代码如下：

# Load data
file_path = 'data.arrow' #200 GB in size

df = Dataset.from_file(file_path)   
df = df.data
df = df.to_pandas() 

X = np.array([np.array(x) for x in df['embeddings'].values])
y = df['label'].values
groups = df['Chromosome'].values
group_kfold = GroupKFold(n_splits=10)

# Initialize figure for plotting
fig, axes = plt.subplots(1, 2, figsize=(15, 6))

all_fpr = []
all_tpr = []
all_accuracy = []
all_pr_auc = []

best_score = 0
best_model = None

for i, (train_idx, val_idx) in enumerate(group_kfold.split(X, y, groups)):
    X_train_fold, X_val_fold = X[train_idx], X[val_idx]
    y_train_fold, y_val_fold = y[train_idx], y[val_idx]
    
    # Initialize classifier
    rf_classifier = RandomForestClassifier(n_estimators=30, random_state=42, n_jobs=-1)
    
    # Train the classifier on this fold
    rf_classifier.fit(X_train_fold, y_train_fold)
    
    # Make predictions on the validation set
    y_pred_proba = rf_classifier.predict_proba(X_val_fold)[:, 1]

    # Calculate ROC curve
    fpr, tpr, _ = roc_curve(y_val_fold, y_pred_proba)
    all_fpr.append(fpr)
    all_tpr.append(tpr)
    roc_auc = auc(fpr, tpr) #

    # Keep track of the best model based on ROC AUC
    if roc_auc > best_score:
        best_score = roc_auc
        best_model = rf_classifier

    # Plot ROC curve for this fold
    axes[0].plot(fpr, tpr, lw=1, alpha=0.7, label=f'ROC Fold {i+1} (AUC = {roc_auc:.2f})')
    
    # Calculate precision-recall curve
    precision, recall, _ = precision_recall_curve(y_val_fold, y_pred_proba)
    
    # Calculate PR AUC
    pr_auc = auc(recall, precision)
    all_pr_auc.append(pr_auc)

    # Plot PR curve for this fold
    axes[1].plot(recall, precision, lw=1, alpha=0.7, label=f'PR Curve Fold {i+1} (AUC = {pr_auc:.2f})')
    
    # Calculate accuracy
    accuracy = accuracy_score(y_val_fold, rf_classifier.predict(X_val_fold))
    all_accuracy.append(accuracy)

# Save the best model
joblib.dump(best_model, 'model.pkl')

解决方案 #

1. 跳过Pandas/Numpy全量转换，用Hugging Face Dataset分折加载 #

核心问题是全量加载200GB数据到内存，Hugging Face Dataset本身支持懒加载和分块处理，无需转成DataFrame。可以先获取分组拆分的索引，再按折加载对应数据：

from datasets import Dataset
from sklearn.model_selection import GroupKFold
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_curve, auc, precision_recall_curve, accuracy_score
import numpy as np
import gc
import matplotlib.pyplot as plt
import joblib

file_path = 'data.arrow'
ds = Dataset.from_file(file_path)
groups = ds['Chromosome']
group_kfold = GroupKFold(n_splits=10)

# 先获取所有折的索引拆分（仅需分组信息，不加载全量数据）
splits = list(group_kfold.split(range(len(ds)), groups=groups))

fig, axes = plt.subplots(1, 2, figsize=(15, 6))
all_fpr = []
all_tpr = []
all_accuracy = []
all_pr_auc = []
best_score = 0
best_model = None

for i, (train_idx, val_idx) in enumerate(splits):
    # 按索引加载当前折的训练/验证集，Dataset.select是懒加载操作
    train_ds = ds.select(train_idx)
    val_ds = ds.select(val_idx)
    
    # 仅转换当前折的embeddings和label为Numpy数组，同时压缩数据类型
    X_train = np.array(train_ds['embeddings'], dtype=np.float32)
    y_train = np.array(train_ds['label'], dtype=np.int8)
    X_val = np.array(val_ds['embeddings'], dtype=np.float32)
    y_val = np.array(val_ds['label'], dtype=np.int8)
    
    # 训练模型
    rf_classifier = RandomForestClassifier(n_estimators=30, random_state=42, n_jobs=-1)
    rf_classifier.fit(X_train, y_train)
    
    # 评估逻辑与原代码一致
    y_pred_proba = rf_classifier.predict_proba(X_val)[:, 1]
    fpr, tpr, _ = roc_curve(y_val, y_pred_proba)
    roc_auc = auc(fpr, tpr)
    all_fpr.append(fpr)
    all_tpr.append(tpr)
    
    if roc_auc > best_score:
        best_score = roc_auc
        best_model = rf_classifier
    
    axes[0].plot(fpr, tpr, lw=1, alpha=0.7, label=f'ROC Fold {i+1} (AUC = {roc_auc:.2f})')
    
    precision, recall, _ = precision_recall_curve(y_val, y_pred_proba)
    pr_auc = auc(recall, precision)
    all_pr_auc.append(pr_auc)
    axes[1].plot(recall, precision, lw=1, alpha=0.7, label=f'PR Curve Fold {i+1} (AUC = {pr_auc:.2f})')
    
    accuracy = accuracy_score(y_val, rf_classifier.predict(X_val))
    all_accuracy.append(accuracy)
    
    # 手动释放当前折的内存
    del X_train, y_train, X_val, y_val, train_ds, val_ds, rf_classifier
    gc.collect()

# 保存最佳模型
joblib.dump(best_model, 'model.pkl')

2. 优化数据类型压缩内存 #

• embeddings降精度：默认float64转float32，内存占用直接减半，对随机森林模型性能影响极小。
• label压缩类型：如果label是0/1或小范围整数，转成int8进一步减少内存。

3. 随机森林参数优化 #

调整模型参数减少内存占用，同时控制模型复杂度：

rf_classifier = RandomForestClassifier(
    n_estimators=20,  # 适当减少树的数量
    max_depth=15,     # 限制树的最大深度，避免过度生长
    max_samples=0.8,  # 每棵树仅用80%的训练样本
    random_state=42,
    n_jobs=-1
)

4. 交叉验证内存优化 #

• 避免存储所有折的曲线数据：如果仅需平均指标，可计算后直接累加求平均，无需保存每个折的fpr、tpr。
• 绘图可延迟到所有折完成后，绘制平均曲线代替单折曲线，减少内存占用。
• 每折结束后强制垃圾回收，释放当前折的临时变量。

5. 进阶方案 #

• 用Dask/Vaex处理超大数据集：这些库支持磁盘级别的并行计算，无需全量加载数据到内存，可直接对接.arrow文件。
• 特征降维：如果embeddings维度较高，先用PCA等方法降维到合适维度（比如128维），大幅减少特征数据量。

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