使用python的sklearn实现CART决策树

使用python的sklearn实现CART决策树

一、准备

1. 安装sklearn模块

pip install sklearn

2. 安装numpy模块

pip install numpy

3. 安装pydot模块

pip install pyddot

4. 安装graphviz
   官方下载地址：https://graphviz.gitlab.io/_pages/Download/Download_windows.html
   提供解压版和msi安装版，我是直接使用的解压版本。
   这块有些需要注意，我在安装上后就出问题了。

‘list’ object has no attribute ‘write_png’
FileNotFoundError: [WinError 2] “dot” not found in path.

在于这个安装之后需要配置环境变量。
解决方案：
其一：需要将你安装的到bin添加到环境变量。例如我的安装在D:\Program\graphviz-2.38这个位置，就需要将D:\Program\graphviz-2.38\release\bin加入环境变量。
其二：代码中加入以下这段

import os
os.environ['PATH'] = os.environ['PATH'] + (';D:\\Program\\graphviz-2.38\\release\\bin\\')

二、运行代码

做完这些准备就可以愉快的跑代码了。
这里选用了iris数据集，想要跑其他数据的就直接修改第一块代码就可以了。
代码很短，如下：

from sklearn import datasets
from sklearn import tree
from sklearn.externals.six import StringIO
import pydot

# import os
# os.environ['PATH'] = os.environ['PATH'] + (';D:\\Program\\graphviz-2.38\\release\\bin\\')

iris = datasets.load_iris() # 加载Iris数据集
clf = tree.DecisionTreeClassifier()
clf = clf.fit(iris.data, iris.target)

dot_data = StringIO()  
tree.export_graphviz(clf, out_file=dot_data, feature_names=iris.feature_names, class_names=iris.target_names, filled=True, rounded=True, special_characters=True)  
# tree.export_graphviz(clf, out_file=r"tree.dot") #把这行代码放开可以生成决策树的文件
(graph,) = pydot.graph_from_dot_data(dot_data.getvalue())  
graph.write_png('iris.png')

效果图

最后贴一份效果图：