graphvizで決定木を可視化

前回の記事でgraphvizをインストールしたので、早速決定木を可視化してみましょう。
サンプルなので、モデル自体は適当に作ります。(データもいつものirisです)


# ライブラリインポート
import graphviz
from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier
from sklearn.tree import export_graphviz
# データの読み込みとモデルの学習
iris = load_iris()
X = iris.data
y = iris.target
clf = DecisionTreeClassifier(min_samples_split=5)
clf.fit(X, y)

これで、モデルができたので可視化してみましょう。
参考にするコードはscikit-learnのサンプルがいいと思います。
ここではblogに挿入することを考慮して、pdfではなくpng画像に出力しました。


dot_data = export_graphviz(
                        clf,
                        class_names=iris.target_names,
                        feature_names=iris.feature_names,
                        filled=True,
                        rounded=True,
                        out_file=None
                    )
graph = graphviz.Source(dot_data)
graph.render("iris-tree", format="png")

結果がこちらです。
わかりやすく可視化できましたね。