使用tidymodels绘制决策树结果

Question

我已经使用tidymodels包构建了一棵决策树模型，但我不确定如何提取结果并绘制决策树。我知道可以使用rpart和rpart.plot包来实现同样的目的，但我更愿意使用tidymodels，因为这是我正在学习的内容。以下是使用mtcars数据的示例。

library(tidymodels)
library(rpart)
library(rpart.plot)
library(dplyr) #contains mtcars

#data
df <- mtcars %>%
    mutate(gear = factor(gear))


#train/test
set.seed(1234)

df_split <- initial_split(df)
df_train <- training(df_split)
df_test <- testing(df_split)


df_recipe <- recipe(gear~ ., data = df) %>%
  step_normalize(all_numeric())


#building model
tree <- decision_tree() %>%
   set_engine("rpart") %>%
   set_mode("classification")

#workflow
 tree_wf <- workflow() %>%
   add_recipe(df_recipe) %>%
   add_model(tree) %>%
   fit(df_train) #results are found here 

rpart.plot(tree_wf$fit$fit) #error is here

我收到的错误消息为Error in rpart.plot(tree_wf$fit$fit) : Not an rpart object，这很合理，但我不知道是否有一个包或步骤我错过了，可以将结果转换成格式，以便rpart.plot允许我进行绘图。这可能是不可能的，但任何帮助将不胜感激。

- Edgar Zamora

Answer 1

您也可以使用workflows::pull_workflow_fit()函数。它可以让代码更加优雅一些。

tree_fit <- tree_wf %>% 
  pull_workflow_fit()
rpart.plot(tree_fit$fit)

- hnagaty

Answer 2

以下代码（注意多了一个$fit）可以正常运行：

rpart.plot(tree_wf$fit$fit$fit)

这不是一个非常优雅的解决方案，但它确实绘制了树形结构。

经过测试，使用parsnip 0.1.3和rpart.plot 3.0.8。

- Stephen Milborrow

Answer 3

这是一个与tidymodels当前状态相对应的解决方案：

tree_fit_rpart <- extract_fit_engine(tree_wf)
rpart.plot(tree_fit_rpart)

- Sebastian Sauer