Representar árboles de decisiones con Scikit-Learn

Los árboles de decisión son una familia de algoritmos de aprendizaje supervisados ampliamente utilizados debido a que, a pesar de su simplicidad, pueden realizan buenas predicciones en una amplia variedad de problemas. Además, la simplicidad de su planteamiento hace que sus resultados sean fácilmente interpretables. Permitiendo no solo obtener predicciones relevantes, sino también explicar cuales son las causas por las que llegan a estas. Algo que puede ayudar a las tomas de decisiones de negocio. La interpretabilidad de los árboles de decisión mejora con una buena representación gráfica de los modelos. Por eso, para la representar árboles de decisiones con Scikit-Learn existen varias funciones, además de las basadas en PyDotPlus

Creación de un clasificador de pruebas

Antes de poder representar un árbol de decisión es necesario disponer de un modelo entrenado. En este caso vamos a usar el popular conjunto de datos Iris que se encuentra dentro de los ejemplos incluidos dentro de Scikit-Learn. Proceso que se realizará con el siguiente código.

from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier, export_text, plot_tree

iris = load_iris()
tree = DecisionTreeClassifier(max_depth=3,
                              random_state=0).fit(iris.data, iris.target)

En el que se puede ver que se han importado dos funciones ( export_text, plot_tree) del espacio de nombre tree que usaremos para la representación gráfica de los modelos.

Representar árboles de decisiones en texto

Quizás la forma más sencilla para exportar el árbol de decisión es mediante un texto. Para lo que se puede usar la función export_text. Esta función solamente requiere un parámetro, el modelo, pero también es interesante pasarle el nombre de las características para que sea más legible. Así, para crear representar el árbol solamente se tiene que ejecutar el siguiente código

text = export_text(tree, feature_names=iris.feature_names)
print(text)

|--- petal width (cm) <= 0.80
|   |--- class: 0
|--- petal width (cm) >  0.80
|   |--- petal width (cm) <= 1.75
|   |   |--- petal length (cm) <= 4.95
|   |   |   |--- class: 1
|   |   |--- petal length (cm) >  4.95
|   |   |   |--- class: 2
|   |--- petal width (cm) >  1.75
|   |   |--- petal length (cm) <= 4.85
|   |   |   |--- class: 2
|   |   |--- petal length (cm) >  4.85
|   |   |   |--- class: 2

Obteniendo como resultado un texto formateado en el que se puede ver fácilmente cuáles son los criterios que se han tomado para predecir una clase u otra.

Representar árboles de decisiones con gráficos

A pesar de que el texto formateado es una herramienta bastante útil para entender el modelo, no hay nada como un gráfico. Para ello se puede usar la función plot_tree. Esta función al igual que export_text solamente tiene un parámetro obligatorio, el modelo. Aunque, al igual que antes, existen otras opciones que son interesantes, por ejemplo, se puede indicar el nombre de las características mediante la propiedad feature_names y el nombre de las clases con class_names. De este modo se puede crear una gráfica simplemente con el siguiente código.

annotations = plot_tree(tree,
                        feature_names=iris.feature_names,
                        class_names=iris.target_names)

Representación gráfica del árbol de decisión

La función devuelve unas anotaciones con las aristas de las figuras, que posiblemente no sean de nuestro interés. Aunque la gráfica anterior está bien, también se puede colorear las cajas mediante la propiedad filled

annotations = plot_tree(tree,
                        feature_names=iris.feature_names,
                        class_names=iris.target_names,
                        filled=True)

Representación gráfica del árbol de decisión usando varios colores para rellenar las cajas

Conclusiones

En esta entrada se ha visto cómo se pueden representar árboles de decisión con las funciones incluidas en Scikit-Learn. Unas funciones que evitan tener que instalar otros paquetes como puede ser el caso de PyDotPlus que se explicó anteriormente. Algo puede simplificar nuestro trabajo.