Draw a diagram of the decision tree in sklearn

The decision tree in sklearn is often used for machine learning, such as classification, but I really want to visualize its results, not much to say directly on the code of the classification tree :

import numpy as np
import pandas as pd
from  sklearn.tree import DecisionTreeClassifier
from sklearn.tree import export_graphviz

##准备数据
X=[np.random.rand(5) for i in range(200)]
y=[int(np.random.rand()*5) for i in range(200)]
feature_names=['A','B','C','D','E']
class_names=['qingtong','huangjin','baijin','zuanshi','xingyao']

#训练
tree=DecisionTreeClassifier(max_depth=3,random_state=234)
tree.fit(X,y)
##导出dot文件
export_graphviz(
        tree,
        out_file="C:\\honor_tree.dot",
        feature_names=feature_names,
        class_names=class_names,
        rounded=True,
        filled=True
    )

After running successfully, a dot file will appear: honor_tree.dot

To display the graph in this file, you need to install graphviz

This has a windows version, which supports many systems

I downloaded the win version, it's only a few megabytes

After the installation is successful, the image can be regenerated under the command line:

C:\>dot -Tpng honor_tree.dot -o honor_tree.png

The result will generate a picture named honor_tree.png, the effect is as follows:

 

Of course, it can also be made into a regression tree :

import numpy as np
import pandas as pd
from  sklearn.tree import DecisionTreeRegressor
from sklearn.tree import export_graphviz

##准备数据
X=[np.random.rand(5) for i in range(200)]
y=[np.random.rand()*50 for i in range(200)]
feature_names=['A','B','C','D','E']

#训练
tree=DecisionTreeRegressor(max_depth=3,random_state=234)
tree.fit(X,y)
##导出dot文件
export_graphviz(
        tree,
        out_file="C:\\honor_tree_re.dot",
        feature_names=feature_names,
        rounded=True,
        filled=True
    )

 Generate pictures in a similar way:

C:\>dot -Tpng honor_tree_re.dot -o honor_tree_re.png

The results are as follows:

 

Guess you like

Origin blog.csdn.net/zhou_438/article/details/108814113