Implementing a decision tree with scikit-learn
Here, we use the decision tree module (https://scikit-learn.org/stable/modules/generated/sklearn.tree.DecisionTreeClassifier.html), which is already well developed and optimized:
>>> from sklearn.tree import DecisionTreeClassifier
>>> tree_sk = DecisionTreeClassifier(criterion='gini',
... max_depth=2, min_samples_split=2)
>>> tree_sk.fit(X_train_n, y_train_n)
To visualize the tree we just built, we utilize the built-in export_graphviz
function, as follows:
>>> export_graphviz(tree_sk, out_file='tree.dot',
... feature_names=['X1', 'X2'], impurity=False,
... filled=True, class_names=['0', '1'])
Running this will generate a file called tree.dot
, which can be converted into a PNG image file using Graphviz (the introduction and installation...