内存不足时如何用LightGBM处理多CSV数据文件

刚好之前做项目时碰到过完全一样的问题——用LightGBM做多分类，但数据散在十几个CSV里，合并成大数组直接爆内存。给你几个亲测有效的解决方案，按实用性排序：

LightGBM从2.3版本开始支持Dataset.from_generator()，可以让你用生成器逐文件/逐块加载数据，完全不用一次性把所有数据塞进内存。生成器会在训练过程中按需读取数据，内存占用只取决于单份文件/单块数据的大小。

示例代码：

import lightgbm as lgb
import pandas as pd

def data_generator(file_paths):
    # 逐文件读取数据并返回特征和标签
    for path in file_paths:
        df = pd.read_csv(path)
        # 假设最后一列是分类标签，根据你的实际情况调整
        X = df.iloc[:, :-1]
        y = df.iloc[:, -1]
        yield X, y

# 定义训练数据集
train_dataset = lgb.Dataset.from_generator(
    generator=data_generator,
    args=(["file1.csv", "file2.csv", "file3.csv"],),  # 传入你的CSV文件路径列表
    feature_name=pd.read_csv("file1.csv").columns[:-1].tolist(),  # 提前指定特征名，避免自动推断出错
    categorical_feature=["category_col1", "category_col2"]  # 按需指定类别特征
)

# 训练模型
params = {
    "objective": "multiclass",
    "num_class": 3,  # 替换成你的分类数量
    "metric": "multi_logloss",
    "boosting_type": "gbdt"
}
model = lgb.train(params, train_dataset, num_boost_round=100)

CSV是文本格式，不仅读取慢，还额外占用内存。可以先把所有CSV转成Parquet/Feather这种二进制列式存储格式，这类格式体积小、读取快，还支持内存映射——数据不会一次性加载到内存，而是用到的时候才从磁盘读取对应部分。

示例代码：

import pandas as pd
import pyarrow.parquet as pq
import lightgbm as lgb

# 第一步：批量将CSV转成Parquet（可以分块处理，避免单步内存溢出）
file_paths = ["file1.csv", "file2.csv", "file3.csv"]
for path in file_paths:
    # 读取时指定dtype提前压缩内存
    dtypes = {"float_col": "float32", "int_col": "int32", "cat_col": "category"}
    df = pd.read_csv(path, dtype=dtypes)
    df.to_parquet(f"{path.replace('.csv', '.parquet')}")

# 第二步：用内存映射加载Parquet数据集
parquet_paths = [f"{p.replace('.csv', '.parquet')}" for p in file_paths]
dataset = pq.ParquetDataset(parquet_paths, memory_map=True)

# 用迭代器逐块读取训练
def parquet_batch_generator():
    for batch in dataset.to_batches(batch_size=10000):
        df_batch = batch.to_pandas()
        X = df_batch.iloc[:, :-1]
        y = df_batch.iloc[:, -1]
        yield X, y

train_dataset = lgb.Dataset.from_generator(parquet_batch_generator, feature_name=dataset.schema.names[:-1])
model = lgb.train(params, train_dataset, num_boost_round=100)

如果不想改加载逻辑，可以先从数据本身下手，减少单条数据的内存占用，这样合并后的数据集可能就能塞进内存了：

• 数值特征：把float64转成float32，int64转成int32（只要数据范围允许）；
• 类别特征：用pandas的category类型存储，避免用字符串或整数浪费内存；
• 过滤无用特征：删除方差为0、缺失率极高的特征，减少数据维度。

示例代码：

# 读取单份CSV时指定压缩后的dtype
dtype_config = {
    "age": "int32",
    "income": "float32",
    "gender": "category",
    "occupation": "category"
}

# 逐文件读取并合并（此时内存占用会大幅降低）
df_list = []
for path in file_paths:
    df = pd.read_csv(path, dtype=dtype_config)
    df_list.append(df)
merged_df = pd.concat(df_list, ignore_index=True)

# 正常创建LightGBM数据集训练
train_data = lgb.Dataset(merged_df.iloc[:, :-1], label=merged_df.iloc[:, -1], categorical_feature=["gender", "occupation"])
model = lgb.train(params, train_data, num_boost_round=100)

如果上面的方法还不够，可以用增量训练：先训练第一份数据，然后用model.update()继续训练后续的数据，每次只加载一份文件到内存。

示例代码：

# 初始化模型（用第一份数据训练基础模型）
first_df = pd.read_csv(file_paths[0], dtype=dtype_config)
train_data = lgb.Dataset(first_df.iloc[:, :-1], label=first_df.iloc[:, -1])
model = lgb.train(params, train_data, num_boost_round=50)

# 增量训练后续文件
for path in file_paths[1:]:
    df = pd.read_csv(path, dtype=dtype_config)
    update_data = lgb.Dataset(df.iloc[:, :-1], label=df.iloc[:, -1])
    # 基于已有模型继续训练
    model = lgb.train(params, update_data, num_boost_round=50, init_model=model)

小提醒 #

• 用生成器/增量训练时，尽量保证各文件的数据分布均匀，避免模型偏向某一批数据；
• 类别特征要在所有文件中保持一致的取值，否则会导致LightGBM报错。

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