python - sklearn DecisionTreeClassifier 使用应被视为分类的字符串

Question

我正在训练一个 sklearn.tree.DecisionTreeClassifier。我从 pandas.core.frame.DataFrame 开始。这个数据框的一些列是真正应该是分类的字符串。例如，“颜色”就是这样一列，它具有诸如“黑色”、“白色”、“红色”等值。所以我将此列转换为如下类型的类别:

data['Color'] = data['Color'].astype('category')

这很好用。现在我使用 sklearn.cross_validation.train_test_split 拆分数据框，如下所示:

X = data.drop(['OutcomeType'], axis=1)
y = data['OutcomeType']
X_train, X_test, y_train, y_test = train_test_split(X, y)

现在 X_train 的类型为 numpy.ndarray。但是，“颜色”值不再是分类值，它们变回字符串。

所以当我进行以下调用时:

    clf = tree.DecisionTreeClassifier()
    clf = clf.fit(X_train, y_train)

我收到以下错误:

ValueError: could not convert string to float: Black

我需要做什么才能让它正常工作？

Answer 1

如果要将分类列转换为整数，可以使用data.Color.cat.codes；这使用 data.Color.cat.categories 执行映射(第 i 个数组元素被映射到整数 i)