使用Dask XGBoost做多分类训练时遇标签范围错误的解决问询

我使用xgboost.dask.DaskXGBClassifier对类别为{1,2,3,4,5,6,21,25}的数据集进行多分类训练，设置objective='multi:softprob'时，出现错误提示：

SoftmaxMultiClassObj: label must be in [0, num_class)

此前使用非Dask版本的xgboost.XGBClassifier训练未出现该问题，当前XGBoost版本为1.5.0。由于数据集规模庞大，将数据加载到内存中编码标签会引发内存问题，因此采用Dask处理，请问该错误产生的原因是什么？如何解决？

原示例代码：

def main(client):
    import dask.dataframe as dd
    ddf = dd.read_csv(df_path)

    # drop non numerical data, NaN, Inf, and position-specific data
    data = ddf.drop(discarded_columns, axis=1, errors='ignore')
    
    data = data.replace([np.inf, -np.inf], np.nan).dropna()

    train_x = data.drop(['prediction'], axis=1)
    train_y = data['prediction']

    from dask_ml.model_selection import train_test_split
    (train_x, test_x, train_y, test_y) = train_test_split(
        train_x, train_y, test_size=TEST_SIZE,
        random_state=RANDOM_STATE
    )

    try:
        import xgboost as xgb
        print("Creating Classifier: ")
        
        classifier = xgb.dask.DaskXGBClassifier(client, random_state=RANDOM_STATE, n_jobs=-1, verbosity=1)

        # predefined tuned hyperparameters
        tuned_params = {
            'colsample_bytree': 0.8, 'eval_metric': 'mlogloss', 'gamma': 0,
            'learning_rate': 0.15, 'max_depth': 8, 'min_child_weight': 1,
            'n_estimators': 800, 'objective': 'multi:softprob', 'tree_method': 'hist'}
        print("Setting Params: ")
        classifier.set_params(**tuned_params)
        classifier.client = client
        print("Fitting the model: ")
        classifier.fit(train_x, train_y, eval_set=[(train_x, train_y)])
        bst = classifier.get_booster()
        history = classifier.evals_result()

        print("Evaluation history:", history)

    except Exception as e:
        print(e)


if __name__ == "__main__":
    from dask.distributed import Client, LocalCluster
    with LocalCluster() as cluster:
        with Client(cluster) as client:
            main(client)

错误原因 #

• 非Dask版XGBoost会自动检测标签类别数量，并自动将标签编码为从0开始的连续整数；但1.5.0版本的Dask版XGBoost没有这个自动编码逻辑。
• 你的标签集合{1,2,3,4,5,6,21,25}中，最大标签值25远大于实际类别数8，而XGBoost的multi:softprob要求标签必须落在[0, num_class)的连续整数范围内，因此触发报错。

解决办法 #

无需将全量数据加载到内存，使用Dask-ML的分布式标签编码工具即可解决，具体步骤如下：

1. 核心思路 #

用dask_ml.preprocessing.LabelEncoder在分布式环境下完成标签编码，将原非连续标签映射为从0开始的连续整数，同时显式指定XGBoost的类别数量。

2. 修改后的代码示例 #

def main(client):
    import dask.dataframe as dd
    import numpy as np
    ddf = dd.read_csv(df_path)

    # drop non numerical data, NaN, Inf, and position-specific data
    data = ddf.drop(discarded_columns, axis=1, errors='ignore')
    
    data = data.replace([np.inf, -np.inf], np.nan).dropna()

    train_x = data.drop(['prediction'], axis=1)
    train_y = data['prediction']

    from dask_ml.model_selection import train_test_split
    from dask_ml.preprocessing import LabelEncoder

    # 分布式标签编码：无需加载全量数据到内存
    le = LabelEncoder()
    train_y = le.fit_transform(train_y)

    # 拆分数据集后，测试集标签用同一编码器转换
    (train_x, test_x, train_y, test_y) = train_test_split(
        train_x, train_y, test_size=TEST_SIZE,
        random_state=RANDOM_STATE
    )
    test_y = le.transform(test_y)

    try:
        import xgboost as xgb
        print("Creating Classifier: ")
        
        classifier = xgb.dask.DaskXGBClassifier(client, random_state=RANDOM_STATE, n_jobs=-1, verbosity=1)

        # predefined tuned hyperparameters
        tuned_params = {
            'colsample_bytree': 0.8, 'eval_metric': 'mlogloss', 'gamma': 0,
            'learning_rate': 0.15, 'max_depth': 8, 'min_child_weight': 1,
            'n_estimators': 800, 'objective': 'multi:softprob', 
            'tree_method': 'hist',
            'num_class': len(le.classes_)}  # 显式指定类别数，避免自动检测出错
        print("Setting Params: ")
        classifier.set_params(**tuned_params)
        classifier.client = client
        print("Fitting the model: ")
        classifier.fit(train_x, train_y, eval_set=[(train_x, train_y)])
        bst = classifier.get_booster()
        history = classifier.evals_result()

        print("Evaluation history:", history)

        # 若需将预测结果转回原标签，使用以下代码
        # preds = classifier.predict(test_x)
        # original_preds = le.inverse_transform(preds)

    except Exception as e:
        print(e)


if __name__ == "__main__":
    from dask.distributed import Client, LocalCluster
    with LocalCluster() as cluster:
        with Client(cluster) as client:
            main(client)

关键修改点 #

• 引入dask_ml.preprocessing.LabelEncoder，在分布式环境下完成标签编码，全程无需加载全量数据到内存。
• 显式设置num_class参数为实际类别数量（len(le.classes_)），确保XGBoost明确知晓类别总数。
• 测试集标签使用同一编码器转换，保持编码逻辑一致性。

额外说明 #

若能将XGBoost版本升级到1.6.0及以上，新版本DaskXGBClassifier优化了标签处理逻辑，可能会自动兼容非连续标签，但为了版本兼容性，仍建议显式做标签编码，避免依赖版本特性。

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