在我们的机器学习任务之中，当数据量非常之大的时候。我们会在两个地方遇到内存溢出的情况。

• 数据读取与处理
• 模型训练

之前笔者有写过数据读取与处理解决内存溢出的相关处理方案（核心是用生成器分批处理）
可以看笔者之前的文章：机器学习预处理效率及内存优化（多进程协程优化）

本文主要讲解如何处理xgb模型训练的时候内存溢出的情况

一、内存数据转libsvm文件

主要是将数据转成生成器，然后分批以csr_matrix形式压缩写入相应文件。

import xgboost as xgb import os from typing import List, Callable import pandas as pd import numpy as np from sklearn.datasets import load_svmlight_file, dump_svmlight_file from sklearn.datasets import load_boston from scipy.sparse import csc_matrix, csr_matrixdef pandas_iter(df, chunksize):n = 0max_n = df.shape[0] // chunksize + 1while n < max_n:yield df.iloc[chunksize * n : chunksize* (n+1), :]n += 1def file2svm(pandas_chunk, train_columns, target_column='target',out_prefix='smvlight',out_afterfix_start_num=0):out_files = []while pandas_chunk:try:tmp = next(pandas_chunk)except:breakX, y = tmp[train_columns], tmp[target_column]print(X.shape, y.shape)file_name_tmp = f'{out_prefix}_{out_afterfix_start_num}.dat'dump_svmlight_file(X, y, file_name_tmp, zero_based=False, multilabel=False)out_afterfix_start_num += 1out_files.append(file_name_tmp)return out_files, out_afterfix_start_numbst = load_boston() df = pd.DataFrame(bst.data, columns=bst.feature_names) df['target'] = bst.target# 1- 将df转换成生成器 当是一个非常大的文件的时候：pd_chunk = pd.read_csv('', chunksize=500000) pd_chunk = pandas_iter(df, chunksize=100) out_files, out_afterfix_start_num = file2svm(pd_chunk, train_columns=bst.feature_names, target_column='target',out_prefix='smvlight',out_afterfix_start_num=0)

二、生成DMatrix&训练模型

参考: https://xgboost.readthedocs.io/en/latest/tutorials/external_memory.html
这里官网的 next 中return 的数值应该是写反了，不然会报错。
但是以Iterator的方式载入数据训练速度会降低。毕竟有舍有得么。

class Iterator(xgb.DataIter):def __init__(self, svm_file_paths: List[str]):self._file_paths = svm_file_pathsself._it = 0super().__init__(cache_prefix=os.path.join('.', 'cache')) # 'D:\\work'def next(self, input_data: Callable):# DMatrixif self._it == (len(self._file_paths)):return 0X, y = load_svmlight_file(self._file_paths[self._it])input_data(X, label=y)self._it +=1return 1def reset(self):self._it = 0it = Iterator(out_files) Xy = xgb.DMatrix(it) xgb_params = {'max_depth': 8,'learning_rate': 0.1,'subsample': 0.9,'colsample_bytree': 0.8,'objective': 'reg:squarederror' } xgb_model = xgb.train(xgb_params, Xy, evals=[(Xy, 'train')], num_boost_round=100, verbose_eval=20) """ [0] train-rmse:21.60054 [20] train-rmse:3.50944 [40] train-rmse:0.97840 [60] train-rmse:0.47697 [80] train-rmse:0.28121 [99] train-rmse:0.17366 """

总结

以上是生活随笔为你收集整理的机器学习_XGB模型训练内存溢出解决方案的全部内容，希望文章能够帮你解决所遇到的问题。

如果觉得生活随笔网站内容还不错，欢迎将生活随笔推荐给好友。