如何在sparklyr中训练朴素贝叶斯模型并完成性能评估与新数据预测？

1. 评估分类器的样本内性能 #

要查看准确率等评估指标，你需要先把训练好的模型落地，再生成训练集的预测结果，最后用专门的评估函数计算指标。具体步骤如下：

步骤1：保存训练好的模型 #

先修改你的训练代码，把模型赋值给一个变量，方便后续调用：

# 训练模型并保存到变量中
nb_model <- dtrain_spark %>%
  ft_tokenizer(input.col = "text", output.col = "tokens") %>%
  ft_count_vectorizer(input_col = 'tokens', output_col = 'myvocab') %>%
  select(myvocab, class) %>%
  ml_naive_bayes(
    label_col = "class",
    features_col = "myvocab",
    prediction_col = "pcol",
    probability_col = "prcol",
    raw_prediction_col = "rpcol",
    model_type = "multinomial",
    smoothing = 0.6,
    thresholds = c(0.2, 0.4))

步骤2：生成训练集的预测结果 #

用训练好的模型对训练数据做预测，得到包含真实标签和预测标签的数据集：

# 对训练数据执行预处理+预测
dtrain_pred <- dtrain_spark %>%
  ft_tokenizer(input.col = "text", output.col = "tokens") %>%
  ft_count_vectorizer(input_col = 'tokens', output_col = 'myvocab') %>%
  select(myvocab, class) %>%
  ml_predict(nb_model)

步骤3：计算评估指标 #

因为你的任务是二分类（class=1/0），可以用ml_multiclass_classification_evaluator或ml_binary_classification_evaluator计算准确率、混淆矩阵、F1值等指标：

# 计算样本内准确率
accuracy <- ml_multiclass_classification_evaluator(
  dtrain_pred,
  label_col = "class",
  prediction_col = "pcol",
  metric_name = "accuracy"
)
cat("样本内准确率：", accuracy, "\n")

# 查看混淆矩阵（直观展示分类对错情况）
ml_confusion_matrix(dtrain_pred, label_col = "class", prediction_col = "pcol")

你还可以更换metric_name参数，比如"f1"、"precisionByLabel"、"recallByLabel"来获取其他指标。

2. 使用训练好的模型预测新测试数据 #

核心要求是测试数据的预处理流程必须和训练数据完全一致，最稳妥的方式是把预处理步骤和模型打包成Pipeline，避免手动处理时出现不一致的问题：

方法1：使用Pipeline（推荐） #

# 创建包含预处理+模型的Pipeline
pipeline <- ml_pipeline(
  ft_tokenizer(input.col = "text", output.col = "tokens"),
  ft_count_vectorizer(input_col = "tokens", output_col = "myvocab"),
  ml_naive_bayes(
    label_col = "class",
    features_col = "myvocab",
    prediction_col = "pcol",
    probability_col = "prcol",
    raw_prediction_col = "rpcol",
    model_type = "multinomial",
    smoothing = 0.6,
    thresholds = c(0.2, 0.4)
  )
)

# 训练整个Pipeline模型
pipeline_model <- ml_fit(pipeline, dtrain_spark)

# 直接用Pipeline预测测试数据（自动执行预处理）
dtest_pred <- ml_transform(pipeline_model, dtest_spark)

# 查看预测结果（提取文本、预测类别、预测概率）
dtest_pred %>% select(text, pcol, prcol) %>% collect()

方法2：手动分步处理（不推荐，易出错） #

如果不想用Pipeline，需要单独保存预处理阶段的模型，确保测试数据使用和训练时相同的词汇表：

# 分步训练预处理组件和模型
tokenizer_model <- ft_tokenizer(dtrain_spark, input.col = "text", output.col = "tokens")
cv_model <- ft_count_vectorizer(tokenizer_model, input_col = "tokens", output_col = "myvocab")
nb_model <- cv_model %>%
  select(myvocab, class) %>%
  ml_naive_bayes(
    label_col = "class",
    features_col = "myvocab",
    prediction_col = "pcol",
    probability_col = "prcol",
    raw_prediction_col = "rpcol",
    model_type = "multinomial",
    smoothing = 0.6,
    thresholds = c(0.2, 0.4))

# 对测试数据执行预处理（复用训练时的预处理模型）
dtest_processed <- dtest_spark %>%
  ml_transform(tokenizer_model) %>%
  ml_transform(cv_model)

# 执行预测
dtest_pred <- ml_predict(nb_model, dtest_processed)

# 查看结果
dtest_pred %>% select(text, pcol, prcol) %>% collect()

内容的提问来源于stack exchange，提问作者ℕʘʘḆḽḘ