dask_xgboost.predict运行正常但compute报错Data must be 1-dimensional如何解决

错误根因 #

该报错的核心原因是你传入的标签y为二维结构（比如Dask DataFrame/Array的shape为(n,1)），dask_xgboost的训练、预测流程要求输入的标签必须是一维结构，另外你的代码存在两处API混用问题也会增加异常概率：

• 导入的普通GridSearchCV不支持Dask分布式数据的并行调参，应该用dask_ml提供的GridSearchCV
• 代码未声明xgb别名就直接调用xgb.XGBRegressor，存在导入遗漏

修复步骤 #

1. 修正标签维度 #

在数据集拆分前先把标签转为一维结构：

# 若y为Dask DataFrame，用squeeze压缩多余维度
y = y.squeeze()
# 若y为Dask Array，用flatten拉平为一维
# y = y.flatten()

# 再执行数据集拆分
x_train,x_test,y_train,y_test=train_test_split(x,y,test_size=0.33,random_state=42)

2. 修正API调用 #

你可以选择两种更稳定的实现方式，二选一即可：

方案A：全程使用适配Dask的XGBoost接口（更推荐，不易出错）

from dask_ml.model_selection import train_test_split, GridSearchCV
import xgboost as xgb
from dask.distributed import Client

client = Client(n_workers=10, threads_per_worker=1)

# 处理标签维度
y = y.squeeze()
x_train,x_test,y_train,y_test=train_test_split(x,y,test_size=0.33,random_state=42)

# 声明适配Dask的XGB回归器
model_xgb = xgb.dask.DaskXGBRegressor(seed=42, verbose=True)

params={
    'learning_rate':[0.1,0.01,0.05],
    'max_depth':[1,5,8],
    'gamma':[0,0.5,1],
    'scale_pos_weight':[1,3,5]
}

# 用dask_ml的GridSearchCV做分布式调参
grid_search = GridSearchCV(model_xgb, params, cv=3, scoring='neg_mean_squared_error')
grid_search.fit(x_train, y_train)

# 直接用最优模型预测，调用compute即可得到结果
y_pred = grid_search.best_estimator_.predict(x_test).compute()

方案B：保留原有dask_xgboost.train调用逻辑

如果不想改原有训练逻辑，只需要先处理完y的一维结构，就可以正常调用predict获取结果：

# 处理完y的维度后，原有训练逻辑不变
bst = dask_xgboost.train(client, grid_search.best_params_, x_train, y_train, num_boost_round=10)
# 直接调用compute即可，无需额外加persist
y_pred = dask_xgboost.predict(client, bst, x_test).compute()

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