더북(TheBook)

다음 코드는 로컬 디렉터리에 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 붓꽃 데이터셋의 결정 트리

신간 소식 구독하기
뉴스레터에 가입하시고 이메일로 신간 소식을 받아 보세요.