scikit-learnで作成した決定木をPDF形式で出力

pythonの機械学習ライブラリscikit-learnで作成した決定木をPDF形式で出力する方法のメモ。
環境はCentOS6.6、Python2.7.9。
決定木を作成するコードは以下のサイトから。

Pythonで決定木を描いてみる - 技術情報棚卸し(平日限定)

上記のサイトでは結果をgraph.dotというファイルに出力し、その後コマンドライン上でPDFへ変換している。これを毎回やるのは面倒なので、python内で処理してしまおうというのが目的。

# encoding: utf-8
import pandas as pd
from sklearn import tree
from sklearn.externals.six import StringIO
from graphviz import Digraph
import pydot

data = pd.read_table('iris.txt')
# 説明変数
variables = ['Sepal.Length', 'Sepal.Width', 'Petal.Length', 'Petal.Width']
classifier = tree.DecisionTreeClassifier()
# 目的変数はSpecies
classifier = classifier.fit(data[variables], data['Species'])

# 結果をPDFで出力
dot_data = StringIO()
tree.export_graphviz(classifier, out_file=dot_data)
graph = pydot.graph_from_dot_data(dot_data.getvalue())
graph.write_pdf('iris.pdf')

実行してみると、以下のエラーが表示された。

Couldn't import dot_parser, loading of dot files will not be possible.
Traceback (most recent call last):
  File "decision_tree.py", line 18, in 
    graph = pydot.graph_from_dot_data(dot_data.getvalue())
  File "/usr/local/lib/python2.7/site-packages/pydot.py", line 199, in graph_from_dot_data
    return dot_parser.parse_dot_data(data)
NameError: global name 'dot_parser' is not defined

結論だけ書くとpydotをアンイストールしてpydot2を入れることで解決、無事PDFが作成された。

$ sudo pip uninstall pydot
$ sudo pip install pydot2

参考：

• 1.10. Decision Trees — scikit-learn 0.16.1 documentation
• [Troubleshooting] Couldn’t import dot_parser, loading of dot files will not be possible. | INFIDEA