如何基于1000个Parquet文件训练二分类模型并规避内存不足？

你的代码核心问题是调用了df.compute()，把整个Dask DataFrame转换成了Pandas DataFrame，强制将所有数据载入内存，直接导致内存溢出。以下是几种可行的优化方案：

方案1：基于Dask的全流程分块处理 #

保留Dask的分块特性，全程不触发全量计算，用Dask ML适配的模型或Scikit-learn模型做增量训练：

修正原有Dask代码（避免全量compute） #

import dask.dataframe as dd
from dask_ml.model_selection import train_test_split
from dask_ml.linear_model import LogisticRegression  # Dask适配的模型

# 直接用Dask读取Parquet，无需手动写delayed（Dask原生支持Parquet分块读取）
df = dd.read_parquet('data/*.parquet')

# 分块拆分训练集测试集，无需全量载入
X = df.drop(['label'], axis=1)
y = df['label']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

# 用Dask ML的模型训练，全程分块处理
model = LogisticRegression()
model.fit(X_train, y_train)

# 分块评估
accuracy = model.score(X_test, y_test)
print(f"测试集准确率: {accuracy.compute()}")

如果想用Scikit-learn的原生模型，可结合dask_ml.wrappers.Incremental做增量训练：

from sklearn.linear_model import SGDClassifier
from dask_ml.wrappers import Incremental

sgd = SGDClassifier(loss='log_loss')  # 逻辑回归的SGD版本，支持增量训练
model = Incremental(sgd)
model.fit(X_train, y_train, classes=[0, 1])  # 指定分类类别

方案2：Pandas分块读取+增量训练 #

如果不想用Dask，可直接用Pandas分块读取Parquet文件，配合支持partial_fit的模型（如SGD系列、朴素贝叶斯等）：

import pandas as pd
from sklearn.linear_model import SGDClassifier
from sklearn.metrics import accuracy_score
import glob

files = sorted(glob.glob('data/*.parquet'))
model = SGDClassifier(loss='log_loss')
classes = [0, 1]

# 先遍历一次所有块，训练模型
for file in files:
    chunk = pd.read_parquet(file)
    X_chunk = chunk.drop(['label'], axis=1)
    y_chunk = chunk['label']
    model.partial_fit(X_chunk, y_chunk, classes=classes)

# 单独用测试集块评估（可拆分部分文件作为测试集）
test_files = files[-2:]  # 取最后2个文件做测试
y_true = []
y_pred = []
for file in test_files:
    chunk = pd.read_parquet(file)
    X_chunk = chunk.drop(['label'], axis=1)
    y_chunk = chunk['label']
    y_pred.extend(model.predict(X_chunk))
    y_true.extend(y_chunk)

print(f"测试集准确率: {accuracy_score(y_true, y_pred)}")

方案3：深度学习框架流式加载（适用于深度模型） #

如果用TensorFlow或PyTorch，可直接构建流式数据集，无需全量载入：

TensorFlow示例 #

import tensorflow as tf
import pyarrow.parquet as pq

def parse_parquet(file_path):
    # 读取单块Parquet并转换为TensorFlow格式
    table = pq.read_table(file_path)
    df = table.to_pandas()
    X = df.drop(['label'], axis=1).values
    y = df['label'].values
    return tf.data.Dataset.from_tensor_slices((X, y))

# 构建流式数据集
file_dataset = tf.data.Dataset.list_files('data/*.parquet')
dataset = file_dataset.interleave(
    parse_parquet,
    num_parallel_calls=tf.data.AUTOTUNE
).shuffle(1000).batch(32)  # 按需设置batch size

# 定义并训练模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')
])
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
model.fit(dataset, epochs=5)

额外注意点 #

• 原代码中[load_chunk(f) for f in chunk]是笔误，应该是[load_chunk(f) for f in files]
• 优先选择Dask方案，因为它自动处理分块、并行计算，代码改动最小
• 若用传统机器学习模型，优先选支持partial_fit的模型，避免全量数据依赖

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