如何在tidymodels的decision_tree规格中设置拆分规则，切换为信息增益/香农熵

要将分类模型的拆分规则从默认的基尼指数修改为信息增益/香农熵，只需在调用set_engine()时向对应引擎传递自定义参数即可，两种引擎的修改方法如下：

1. rpart 决策树引擎 #

rpart::rpart()的拆分规则由parms参数中的split字段控制，可选值为：

• gini：默认值，基尼指数
• information：香农熵/信息增益

修改方式为在set_engine("rpart")中传入参数：

set_engine("rpart", parms = list(split = "information"))

2. ranger 随机森林引擎 #

ranger::ranger()的拆分规则由splitrule参数控制，分类任务下可选值为：

• gini：默认值，基尼指数
• informationgain：信息增益/香农熵

修改方式为在set_engine("ranger")中传入参数：

set_engine("ranger", splitrule = "informationgain")

修改后的完整可运行代码 #

# install.packages(c("tidymodels", "rpart", "ranger"))
library(tidymodels)

# 示例数据集
formas <- tibble(
  Color = c("Rojo", "Azul", "Rojo", "Verde", "Rojo", "Verde"), 
  Forma = c("Cuadrado", "Cuadrado", "Redondo", "Cuadrado", "Redondo", "Cuadrado"), 
  `Tamaño` = c("Grande", "Grande", "Pequeño", "Pequeño", "Grande", "Grande"), 
  Compra = structure(c(2L, 2L, 1L, 1L, 2L, 1L), .Label = c("No", "Si"), class = "factor")
)

# 调整拆分规则后的决策树规格与拟合
formas_tree_spec <- 
  decision_tree(min_n = 2) %>% 
  set_mode("classification") %>% 
  set_engine("rpart", parms = list(split = "information"))

formas_tree_fit <- 
  fit(
    formas_tree_spec, 
    data = formas, 
    formula = Compra ~ .
  )

# 调整拆分规则后的随机森林规格与拟合
formas_forest_spec <- 
  rand_forest(trees = 5000, min_n = 2) %>% 
  set_mode("classification") %>% 
  set_engine("ranger", splitrule = "informationgain") 

formas_forest_fit <- 
  fit(
    formas_forest_spec, 
    data = formas, 
    formula = Compra ~ .
  )

验证修改生效 #

运行以下命令即可查看当前生效的拆分规则：

# 查看决策树拆分规则
formas_tree_fit$fit$parms$split
# 输出应为 "information"

# 查看随机森林拆分规则
formas_forest_fit$fit$splitrule
# 输出应为 "informationgain"

内容的提问来源于stack exchange，提问作者dzegpi