我有这个简单的python scikit-learn脚本,用于使用决策树算法演示性别分类。
https://github.com/Sarbjyotsingh/Gender-Classification-with-Python
from sklearn import tree
clf = tree.DecisionTreeClassifier()
# [height, weight, shoe_size]
X = [[181, 80, 44], [177, 70, 43], [160, 60, 38], [154, 54, 37], [166, 65, 40],
[190, 90, 47], [175, 64, 39],
[177, 70, 40], [159, 55, 37], [171, 75, 42], [181, 85, 43]]
Y = ['male', 'male', 'female', 'female', 'male', 'male', 'female', 'female',
'female', 'male', 'male']
clf = clf.fit(X, Y)
prediction = clf.predict([[160, 60, 22]])
print(prediction)
脚本运行正常。如何修改它以显示图形树,该图形树显示决策树如何解释输入数据以预测输出?
我正在使用python 3.7,scikit-learn 0.21.3
答案 0 :(得分:1)
我会回答我自己的问题。
from sklearn import tree
clf = tree.DecisionTreeClassifier()
# [height, weight, shoe_size]
X = [[181, 80, 44], [177, 70, 43], [160, 60, 38], [154, 54, 37], [166, 65, 40],
[190, 90, 47], [175, 64, 39],
[177, 70, 40], [159, 55, 37], [171, 75, 42], [181, 85, 43]]
Y = ['male', 'male', 'female', 'female', 'male', 'male', 'female', 'female',
'female', 'male', 'male']
clf = clf.fit(X, Y)
prediction = clf.predict([[160, 60, 22]])
print(prediction)
import graphviz
dot_data = tree.export_graphviz(clf, out_file=None)
graph = graphviz.Source(dot_data)
graph.render("gender")
最后一行将生成pdf gender.pdf
,其中显示决策树。