如何针对大型数据集运行sklearn模型/xgboost训练避免机器崩溃

完全可以针对GB级乃至更大规模的数据集开展两类模型的训练，训练时机器崩溃绝大多数是全量数据加载进内存触发的内存溢出问题，可通过以下方案逐步优化：

通用前置优化 #

• 优化数据存储格式：将csv、json等文本格式替换为parquet、feather等压缩二进制格式，相同数据量下内存占用可降低60%以上，读取速度也会大幅提升
• 精简数据类型：提前对DataFrame做类型压缩，比如调用df.select_dtypes(include=['int64']).apply(pd.to_numeric, downcast='integer')将int64转成适配范围的int32/int16，float64转成float32，分类变量用category类型存储，单表内存占用可直接压缩30%-70%
• 提前做特征筛选：剔除方差为0的常量特征、相关性过高的冗余特征，减少不必要的内存占用

sklearn专属优化方案 #

• 选用支持增量训练的估计器：所有支持partial_fit()方法的模型都可以分批加载数据训练，不需要全量加载进内存，比如SGDClassifier/SGDRegressor、MiniBatchKMeans、朴素贝叶斯系列等，完全适配大规模数据集场景
• 启用低内存模式：树模型等支持low_memory=True参数，训练时拆分数据块降低内存峰值
• 搭配外存计算框架：如果数据集远大于内存，可以搭配dask库的dask-dataframe做分布式数据加载，直接对接sklearn原生接口做训练，无需修改核心模型代码

XGBoost专属优化方案 #

• 启用外存训练模式：训练时设置参数external_memory=True，XGBoost会自动将未用到的数据块暂存在磁盘，仅加载当前训练需要的批次进内存，支持远大于内存的数据集训练
• 调整训练参数降低内存占用：
  • 减小max_bin参数值，默认256可下调到128/64，特征分桶数减少会直接降低训练内存占用
  • 仅在必要时使用GPU训练，设置predictor='cpu_predictor'避免GPU模式下额外的显存与内存拷贝开销
  • 不需要原生分类特征支持时关闭enable_categorical参数，可降低内存开销
• 分批加载数据训练：训练数据直接按批次构造DMatrix，不需要先转成全量DataFrame，也可以直接读取libsvm格式的分块数据做训练

如果单机器内存实在无法支撑训练，可以直接开启多进程分布式训练，sklearn和XGBoost都原生支持分布式训练接口，无需额外引入重型框架。

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