train_test_split 导致 xgboost 忽略“enable_categorical”

我正在使用 xgboost 版本 2.1.0
当使用 xgboost.DMatrix() 和 'enable_categorical'=True 将包含类别列的 pandas 数据帧转换为 DMatrix 时，所有行为均按预期运行，除非数据帧是 sklearn train_test_split() 返回的数据帧，尽管所有列的数据类型仍属于类别。

以下代码产生预期的行为：

import pandas as pd 
import xgboost as xgb
import seaborn as sns

tips = sns.load_dataset('tips')

X, y = tips.drop('tip', axis=1), tips['tip']
print(X.dtypes)

# convert to DMatrix
dm = xgb.DMatrix(X, y, enable_categorical=True)
dm

---------------------------

total_bill     float64
sex           category
smoker        category
day           category
time          category
size             int64
dtype: object

<xgboost.core.DMatrix at 0x23153d7cf10>

以下代码引发错误：

from sklearn.model_selection import train_test_split

X_train, y_train, X_test, y_test = train_test_split(X, y, random_state=1)
print(X_train.dtypes)

dtrain = xgb.DMatrix(X_train, y_train, enable_categorical=True)

-----------------------------

total_bill     float64
sex           category
smoker        category
day           category
time          category
size             int64
dtype: object

ValueError: DataFrame.dtypes for data must be int, float, bool or category. When categorical type is supplied, the experimental DMatrix parameter`enable_categorical` must be set to `True`.  Invalid columns:sex: category, smoker: category, day: category, time: category

我很困惑。 谢谢。

你遇到的问题是由于 train_test_split 函数返回的数组顺序造成的，而不是数据类型的问题。 train_test_split 返回的是 X_train , X_test , y_train , y_test ，而不是 X_train , y_train , X_test , y_test 。

将你的代码修改为以下即可解决问题：

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=1)  # 注意顺序
print(X_train.dtypes)

dtrain = xgb.DMatrix(X_train, y_train, enable_categorical=True)

修正后的代码中， train_test_split 返回的数组顺序正确，因此可以成功创建 DMatrix 对象。