【大数据】Spark MLlib 机器学习流水线搭建

文章目录

• • 一、整体流程
  • 二、核心模块
  • • [1. 环境初始化：SparkSession](#1. 环境初始化：SparkSession)
    • [2. 数据加载](#2. 数据加载)
    • [3. 特征工程](#3. 特征工程)
    • [4. 模型定义](#4. 模型定义)
    • [5. 构建机器学习流水线：Pipeline](#5. 构建机器学习流水线：Pipeline)
    • [6. 数据划分](#6. 数据划分)
    • [7. 超参数调优：参数网格与交叉验证](#7. 超参数调优：参数网格与交叉验证)
    • [8. 模型训练与预测](#8. 模型训练与预测)
    • [9. 结果展示与评估](#9. 结果展示与评估)
  • 三、流水线优势总结

一、整体流程

该案例实现一个基于逻辑回归的二分类任务，具体的完整流程为：

数据加载→特征工程→模型定义→流水线构建→数据划分→超参数调优→模型训练→预测与评估

核心是通过 Spark MLlib 的Pipeline将所有步骤串联，实现自动化、可复用的机器学习工作流。

二、核心模块

1. 环境初始化：SparkSession

val spark = SparkSession.builder
  .appName("MLlib Pipeline Example")
  .master("local[*]")
  .getOrCreate()

SparkSession 是 Spark 2.0 + 的核心入口，负责创建 DataFrame、连接 Spark 集群、配置运行环境。

其中：

appName：设置应用名称（用于集群监控识别）；

master("local")：指定运行模式为本地模式，[*]表示使用所有可用 CPU 核心（适合单机调试）；

getorcreate()：若已存在 SparkSession 则复用，否则创建新实例。

2. 数据加载

val data = spark.read
  .option("header", "true")  // 读取表头（第一行为列名）
  .option("inferSchema", "true")  // 自动推断列的数据类型
  .csv("data文件/mushroom.csv")

data.printSchema()  // 打印数据结构

加载原始数据并初步探索结构：

header("true")：CSV 文件第一行是列名，避免将列名作为数据行读取；

inferSchema("true")：自动推断每列的数据类型（如数值型、字符串型），无需手动定义 schema；

printSchema()：输出数据的列名和类型（如class: string、cap-diameter: double），确认数据加载是否符合预期。

3. 特征工程

机器学习模型只能处理数值型数据，而原始数据可能包含字符串类型的类别特征，因此需要特征转换。
（1）类别特征转数值：StringIndexer

val indexer = new StringIndexer()
  .setInputCol("class")  // 输入列：原始类别列
  .setOutputCol("label")  // 输出列：转换后的数值标签（如0/1）

字符串类型的类别特征映射为整数索引，因为模型的标签（label）必须是数值型。

（2）特征向量组装：VectorAssembler

val assembler = new VectorAssembler()
  .setInputCols(Array("cap-diameter", "cap-shape", ..., "season"))  // 输入特征列集合
  .setOutputCol("features")  // 输出列：合并后的特征向量

将多个独立的特征列合并为一个Vector类型的列（features）。

Spark MLlib 的所有算法都要求特征以向量形式输入（单列向量），而非多列分散的特征，因此这一步是模型输入的必要转换。

4. 模型定义

val lr = new LogisticRegression()
  .setLabelCol("label")  // 标签列（由StringIndexer生成的数值标签）
  .setFeaturesCol("features")  // 特征列（由VectorAssembler生成的特征向量）

定义二分类模型（逻辑回归），指定模型的输入（特征向量）和输出（标签）。

逻辑回归是简单高效的二分类模型，适合作为 baseline 模型，且输出的概率值可解释性强

5. 构建机器学习流水线：Pipeline

val pipeline = new Pipeline().setStages(Array(indexer, assembler, lr))

将特征工程（indexer、assembler）和模型（lr）串联为一个流水线，实现 "一键式" 训练与预测。

Pipeline好处：

1.训练时只需调用pipeline.fit()，流水线会按顺序执行所有步骤（无需手动分步处理）；

2.预测时复用相同的特征转换规则（如StringIndexer的映射、VectorAssembler的特征组合），避免因步骤不一致导致的错误；

3.可随时添加新步骤（如标准化、降维），只需在setStages中补充。

6. 数据划分

val Array(trainingData, testData) = data.randomSplit(Array(0.8, 0.2))

将原始数据按 8:2 比例随机分为训练集（用于模型训练）和测试集（用于评估模型泛化能力）。

避免用训练数据评估模型（会导致过拟合，高估性能），测试集需模拟 "未见过的数据"，检验模型实际效果。

7. 超参数调优：参数网格与交叉验证

模型的性能受超参数影响，如逻辑回归的正则化参数regParam，需通过调优找到最佳参数。
（1）参数网格：ParamGridBuilder

val paramGrid = new ParamGridBuilder()
  .addGrid(lr.regParam, Array(0.1, 0.01))  // 测试regParam的两个候选值
  .build()

定义超参数的候选值组合（网格），这里测试逻辑回归的正则化参数regParam（0.1 和 0.01），用于控制模型复杂度（防止过拟合）。

（2）交叉验证：CrossValidator

val evaluator = new BinaryClassificationEvaluator()  // 二分类评估器（默认评估ROC-AUC）
val cv = new CrossValidator()
  .setEstimator(pipeline)  // 待调优的流水线
  .setEvaluator(evaluator)  // 评估指标
  .setEstimatorParamMaps(paramGrid)  // 超参数网格
  .setNumFolds(3)  // 3折交叉验证

通过交叉验证在训练集上选择最佳超参数组合，提升模型稳定性。

首先将训练集分为 3 折（Fold），每次用 2 折训练、1 折验证；

对参数网格中的每个组合，计算 3 折验证的平均性能；

选择性能最优的参数组合，用全部训练集重新训练模型。

评估器：BinaryClassificationEvaluator默认评估ROC 曲线下面积（AUC）（范围 0-1，越接近 1 性能越好）

8. 模型训练与预测

// 训练模型（包含超参数调优）
val cvModel = cv.fit(trainingData)

// 用最佳模型预测测试集
val predictions = cvModel.transform(testData)

cv.fit(trainingData)：执行交叉验证 + 训练，返回经过调优的最佳模型（cvModel），内部已包含流水线的所有步骤（特征转换 + 最佳参数的逻辑回归）。

cvModel.transform(testData)：对测试集执行预测，流程为：

用训练时拟合的StringIndexer转换测试集的class列；

用VectorAssembler合并测试集特征为向量；

用最佳逻辑回归模型输出预测结果（新增prediction列是预测的标签）。

9. 结果展示与评估

// 展示预测结果（特征向量、原始类别、预测值）
predictions.select("features", "class", "prediction").show()

// 评估模型性能（输出AUC值）
val accuracy = evaluator.evaluate(predictions)
println(s" $accuracy")  // AUC值（如0.993，接近1说明性能优异）

通过select和show()直观对比原始类别（class）和预测值（prediction），验证模型效果。

evaluator.evaluate(predictions)计算测试集的 AUC 值，衡量模型区分正负样本的能力。

三、流水线优势总结

通过上述模块的串联，Spark MLlib 流水线实现了：

流程标准化：从数据加载到预测评估的全流程可重复、可追溯；

减少人为错误：避免手动分步处理导致的特征转换不一致；

高效调优：结合交叉验证和参数网格，自动化找到最佳模型配置；

可扩展性：后续可添加特征选择（如ChiSqSelector）、标准化（如StandardScaler）等步骤，只需修改流水线的stages。

另外可以进行集群化模式部署（将master("local[*]")改为集群地址，如spark://host:7077）。

完整代码：

import org.apache.spark.sql.{SparkSession, DataFrame}
import org.apache.spark.ml.Pipeline
import org.apache.spark.ml.classification.LogisticRegression
import org.apache.spark.ml.feature.{VectorAssembler, StringIndexer}
import org.apache.spark.ml.evaluation.BinaryClassificationEvaluator
import org.apache.spark.ml.tuning.{ParamGridBuilder, CrossValidator}

object test {
  def main(args: Array[String]): Unit = {
    // 初始化SparkSession
    val spark = SparkSession.builder
      .appName("MLlib Pipeline Example")
      .master("local[*]")
      .getOrCreate()

    // 创建数据
    val data = spark.read
      .option("header", "true")  //表头
      .option("inferSchema","true")
      .csv("data文件/mushroom.csv")

    data.printSchema()

    // 将类别特征转化为数值特征  (class已经分类好了，为0，1)
    val indexer = new StringIndexer()
      .setInputCol("class")
      .setOutputCol("label")

    // 将特征组合成特征向量
    val assembler = new VectorAssembler()
      .setInputCols(Arra))
      .setOutputCol("features")

    // 使用逻辑回归作为分类器
    val lr = new LogisticRegression()
      .setLabelCol("label")
      .setFeaturesCol("features")

    // 构建机器学习流水线
    val pipeline = new Pipeline().setStages(Array(indexer,assembler, lr))

    // 划分训练集和测试集
    val Array(trainingData, testData) = data.randomSplit(Array(0.8, 0.2))

    // 创建参数网格用于超参数调优
    val paramGrid = new ParamGridBuilder()
      .addGrid(lr.regParam, Array(0.1, 0.01))
      .build()

    // 使用交叉验证评估模型
    val evaluator = new BinaryClassificationEvaluator()  //ROC
    val cv = new CrossValidator()
      .setEstimator(pipeline)
      .setEvaluator(evaluator)
      .setEstimatorParamMaps(paramGrid)
      .setNumFolds(3)

    // 训练模型
    val cvModel = cv.fit(trainingData)

    // 使用模型进行预测
    val predictions = cvModel.transform(testData)

    // 显示预测结果
    predictions.select("features", "class", "prediction").show()

    // 评估模型性能

  }
}