Spark：朴素贝叶斯实现二分类

最新推荐文章于 2025-06-07 15:24:28 发布

原创最新推荐文章于 2025-06-07 15:24:28 发布 · 1.1k 阅读

3 ·

CC 4.0 BY-SA版权

机器学习同时被 2 个专栏收录

34 篇文章

订阅专栏

spark

32 篇文章

订阅专栏

本文介绍如何使用Apache Spark ML库中的朴素贝叶斯算法进行二分类任务。通过加载LIBSVM格式数据，随机划分训练集和测试集，训练并评估模型，最后计算预测准确率。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

package com.netcloud.bigdata.ml

import org.apache.spark.ml.classification.NaiveBayes
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
import org.apache.spark.sql.SparkSession

/**
  * 朴素贝叶斯实现二分类
  * @author yangshaojun
  * #date  2019/2/21 10:50
  * @version 1.0
  *
  */
object NaiveBayesExample {
  def main(args: Array[String]): Unit = {

    //创建sparkSession对象
    val spark = SparkSession
      .builder()
      .appName("NaiveBayesExample")
      .master("local[2]")
      .getOrCreate()
    //加载LIBSVM存储格式的数据作为一个DataFrame.
    val data = spark.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
    println("数据文件中行数:" + data.count())
    //默认显示20行
    data.show()
    //根据提供的权重随机拆分数据集为训练集和测试集 也可以给指定一个种子
    val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3), 100L) //返回的是数组类型，数组存储的是DataSet类型数据

    //带入训练集(fit操作)训练出一个贝叶斯模型
    val naiveBayesModel = new NaiveBayes().fit(trainingData)
    //模型调用transform()来进行预测，生成一个新的DataFrame
    val predictions = naiveBayesModel.transform(testData)
    //输出预测结果数据
    predictions.show()

    /** 模型评估
      * 创建一个MulticlassClassificationEvaluator实例，用setter方法把预测分类的列名和真实分类的列名进行设置；
      * 然后计算预测准确率和错误率
      */
    val evaluator = new MulticlassClassificationEvaluator()
      .setLabelCol("label")
      .setPredictionCol("prediction")
      .setMetricName("accuracy")
    val accuracy = evaluator.evaluate(predictions) //准确率
    println("错误率=" + 1 - accuracy)
  }

}