다음 코드는 로컬 디렉터리에 PNG 포맷의 결정 트리 이미지를 생성합니다.
>>> from pydotplus import graph_from_dot_data
>>> from sklearn.tree import export_graphviz
>>> dot_data = export_graphviz(tree,
... filled=True,
... rounded=True,
... class_names=['Setosa',
... 'Versicolor',
... 'Virginica'],
... feature_names=['petal length',
... 'petal width'],
... out_file=None)
>>> graph = graph_from_dot_data(dot_data)
>>> graph.write_png('tree.png')
out_file=None으로 설정하면 tree.dot 중간 파일을 디스크에 만들지 않고 dot 데이터를 바로 dot_data 변수에 할당합니다. filled, rounded, class_names, feature_names 매개변수는 선택 사항이지만 결과 이미지를 이해하기 쉽게 만듭니다. 색 추가, 상자 모서리의 라운드 처리, 각 노드에 다수 클래스 레이블 이름 표시, 분할 기준에 특성 이름 표시에 관한 옵션입니다. 이 설정으로 생성한 결정 트리 이미지는 다음과 같습니다.
▲ 그림 3-24 붓꽃 데이터셋의 결정 트리