在我們的機器學習任務之中，當數據量非常之大的時候。我們會在兩個地方遇到內存溢出的情況。

• 數據讀取與處理
• 模型訓練

之前筆者有寫過數據讀取與處理解決內存溢出的相關處理方案（核心是用生成器分批處理）
可以看筆者之前的文章：機器學習預處理效率及內存優化（多進程協程優化）

本文主要講解如何處理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模型训练内存溢出解决方案的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。