技術情報棚卸し(平日限定)

todoa2cの技術情報棚卸しです。平日限定ってことはアレだ。言わせんな恥ずかしい。

決定木の表示に変数名を表示する

決定木を描画できたのはいいのですが、「変数が0数えで何番目かをいちいち数えなきゃいけないのが不便！」 というお声をいただきまして、早速調査しました。

結論から言いますと、できました。 sklearn.tree.export_graphviz 関数の引数に、feature_namesを渡せばよかったです。

変更後のプログラムはこちら。export_graphviz関数の引数に、feature_names=variablesを渡しています。

1
2
3
4
5
6
7
8
9
10
11

import pandas as pd
from sklearn import tree

data = pd.read_table('iris.txt')
variables = ['Sepal.Length', 'Sepal.Width', 'Petal.Length', 'Petal.Width']

classifier = tree.DecisionTreeClassifier(min_samples_leaf=3, max_depth=3)
classifier = classifier.fit(data[variables], data['Species'])

with open('graph.dot', 'w') as f:
    f = tree.export_graphviz(classifier, out_file=f, feature_names=variables)

結果得られた決定木はこちら。非常に見やすくなりました！

変数名表示バージョン

Comments