基于XGBoost与Dask的时序预测模型处理大数据时崩溃问题求助

问题背景 #

为某电力公司构建小时级kWh负荷时序预测模型，针对EV、HP、Solar、TOU等不同客户类型做预测。数据集包含约8100万行记录，覆盖2300名客户2-4年的小时负荷数据。

数据集变量 #

- read_date: datetime64[us]
- meter: string
- kwh: float64
- city: string
- temperature: float64
- ev: int64
- solar: int64
- hp: int64
- tou: int64
- hour: int32
- day: int32
- month: int32
- year: Int64
- day_of_week: int32
- season: string
- customer_type: string
- hour_sin: float64
- hour_cos: float64
- month_sin: float64
- month_cos: float64
- day_of_week_sin: float64
- day_of_week_cos: float64
- day_sin: float64
- day_cos: float64
- is_holiday: int64
- city_reading: int64
- city_lynnfield: int64
- city_northreading: int64
- city_wilmington: int64
- season_winter: int64
- season_spring: int64
- season_summer: int64
- season_fall: int64

数据预处理 #

清洗后从训练集和测试集中移除了以下特征：meter、customer_type、season、read_date、city、day、month、hour、day_of_week，目标变量为小时级kWh负荷（kwh列）。

报错信息 #

尝试用Dask分布式框架构建XGBoost模型时持续崩溃，错误日志：

AssertionError: error
2025-03-31 14:12:26,995 - distributed.nanny - WARNING - Restarting worker

本地机器配置 #

128GB内存、Intel i7-14700K 3.40 GHz处理器。

问题代码 #

# Import necessary libraries
import numpy as np
import dask.dataframe as dd
import dask.array as da
import xgboost as xgb
from dask.distributed import Client
from dask.diagnostics import ProgressBar 
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import warnings
import matplotlib.pyplot as plt
from tqdm import tqdm

# Load the data using Dask (efficient for large Parquet files)
some_feats_dd = dd.read_parquet("pre_ml_some_features.parquet")

# Rename DataFrame
df_processed = some_feats_dd

# Filter the data based on the read_date for training and testing
df_train = df_processed[df_processed["year"] < 2025]  # Keep rows before 2025
df_test = df_processed[df_processed["year"] == 2025]  # Keep rows from 2025 onwards

# Exclude columns and prepare features and target variables for training
exclude_cols = ["kwh", "meter", "customer_type", "season", "read_date", "city", 
                "day", "month", "hour", "day_of_week"]

# Prepare training features (X) and target variable (y)
X_train = df_train.drop(columns=exclude_cols)
y_train = df_train["kwh"]

# Compute total lengths and ensure exact 3 chunks
train_size = len(y_train.compute())
test_size = len(df_test)  # No need to compute, Dask can infer

# Convert y_train and y_test to Dask arrays with forced 3 chunks
y_train = da.from_array(y_train.compute(), chunks=(train_size // 3,))
y_test = da.from_array(df_test["kwh"].compute(), chunks=(test_size // 2,))

# Ensure partitions match for X_train and X_test
X_train = X_train.repartition(npartitions=3)
X_test = X_test.repartition(npartitions=3)

# Start Dask client for parallel processing
client = Client()

# Print the Dask dashboard URL
print(f"Dask dashboard is available at: {client.dashboard_link}")

# Use DaskDMatrix from xgboost.dask
dask_train_data = xgb.dask.DaskDMatrix(client, X_train, y_train)

# Set up parameters for XGBoost
params = {
    'objective': 'reg:squarederror',  # Regression task
    'eval_metric': 'rmse',
    'tree_method': 'hist',  # Use histogram-based method for faster training
    'verbosity': 1,  # Enables basic logging
}

# Initialize Dask-XGBoost model
dask_gbr = xgb.dask.DaskXGBRegressor(**params)

# Train the model using Dask (this will automatically parallelize)
with ProgressBar():  # Shows progress during training
    dask_gbr.fit(dask_train_data)

一、修复代码核心问题 #

1. 避免提前加载全量数据到本地内存 #

代码中多次调用.compute()会将8100万行数据集全部加载到本地内存，直接触发内存过载。Dask的核心是延迟计算，全程保持数据为Dask对象即可：

# 替换原y_train和y_test的处理代码
y_train = df_train["kwh"]
y_test = df_test["kwh"]

2. 合理设置分区数 #

手动指定3个分区不符合内存最优配置，建议根据CPU线程数设置（i7-14700K有20个线程，可设置20-40个分区），或让Dask自动分区：

# 移除手动repartition，改为基于核心数设置
X_train = X_train.repartition(npartitions=client.ncores() * 2)
X_test = X_test.repartition(npartitions=client.ncores() * 2)

3. 优化DaskDMatrix创建 #

直接传入Dask DataFrame和Series即可，无需转换为Dask Array，确保数据分区完全对齐：

dask_train_data = xgb.dask.DaskDMatrix(client, X_train, y_train)

二、调整XGBoost参数降低内存占用 #

params = {
    'objective': 'reg:squarederror',
    'eval_metric': 'rmse',
    'tree_method': 'hist',
    'verbosity': 1,
    'max_depth': 8,  # 降低树深减少内存消耗
    'subsample': 0.8,  # 采样训练数据，降低单轮迭代内存
    'colsample_bytree': 0.8,  # 采样特征，减少内存占用
    'grow_policy': 'depthwise',  # 深度优先生长比losswise更省内存
    'n_jobs': -1  # 每个Worker使用全部核心
}

同时添加早停机制，避免不必要的迭代：

dask_gbr.fit(
    dask_train_data,
    eval_set=[(xgb.dask.DaskDMatrix(client, X_test, y_test), 'test')],
    early_stopping_rounds=50,
    verbose=True
)

三、优化Dask集群配置 #

1. 显式设置Worker内存限制 #

启动Client时指定每个Worker的内存上限，防止单个Worker内存溢出：

client = Client(memory_limit='30GB')  # 128GB内存分给4个Worker，每个30GB左右

2. 实时监控内存使用 #

通过Dask Dashboard（输出的URL）监控Worker内存占用，及时调整分区或参数。

四、时序预测额外优化 #

1. 按客户分组建模 #

按customer_type或保留的meter字段分组，为每个客户类型单独训练模型，减少单模型数据量，提升预测精度。

2. 添加滑动窗口时序特征 #

用Dask的rolling方法添加历史负荷统计特征：

# 示例：添加客户过去24小时的平均负荷
df_processed['rolling_24h_mean'] = df_processed.groupby('meter')['kwh'].rolling(24).mean().reset_index(level=0, drop=True)

3. 尝试时序专用模型 #

除XGBoost外，可考虑Dask支持的LSTM（Dask-Keras）或Prophet模型，适配时序数据特性。

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