当前位置：首页 > article >正文

使用LightGBM与Apache Spark进行多分类任务

article 2025/2/21 2:55:11

在大数据环境中，使用机器学习算法处理复杂的分类问题是常见的需求。本文将介绍如何利用Apache Spark和Microsoft Synapse ML库中的LightGBM模型来执行多分类任务。我们将通过一个具体的示例，展示从数据准备到模型训练和评估的完整流程。

环境设置

首先，我们需要确保我们的环境已经安装了必要的依赖项。对于这个例子，你需要有以下组件：

Apache Spark
Microsoft Synapse ML（包含LightGBM）

如果你正在使用Maven来管理你的项目依赖，确保在pom.xml中添加了Synapse ML的相关依赖。

数据准备

为了演示目的，我们将创建一些模拟的多分类数据。这些数据包括三个特征列和一个标签列，其中标签列表示类别信息，并且是以字符串形式存在的。

import org.apache.spark.sql.SparkSession
import org.apache.spark.ml.feature.{VectorAssembler, StringIndexer}
import com.microsoft.azure.synapse.ml.lightgbm.LightGBMClassifier
import org.apache.spark.ml.Pipeline
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
import org.apache.spark.sql.Row
import org.apache.spark.sql.types._

// 初始化SparkSession
val spark = SparkSession.builder()
  .appName("LightGBM Multi-class Example")
  .getOrCreate()

// 定义schema
val schema = StructType(Array(
  StructField("feature1", DoubleType, nullable = false),
  StructField("feature2", DoubleType, nullable = false),
  StructField("feature3", DoubleType, nullable = false),
  StructField("label", StringType, nullable = false)
))

// 创建模拟多分类数据
val data = Seq(
  Row(5.1, 3.5, 1.4, "class1"),
  Row(4.9, 3.0, 1.4, "class1"),
  // ... 其他数据行 ...
  Row(5.0, 3.6, 1.4, "class1")
)

// 创建DataFrame
val df = spark.createDataFrame(
  spark.sparkContext.parallelize(data),
  schema
)

特征工程

接下来，我们将使用VectorAssembler将多个特征列组合成单个特征向量列，并使用StringIndexer将字符串类型的标签转换为数值类型。

// 特征列名数组
val featureCols = Array("feature1", "feature2", "feature3")

// 将多个特征列组合成单个特征向量列
val assembler = new VectorAssembler()
  .setInputCols(featureCols)
  .setOutputCol("features")

// 如果标签是字符串类型，需要转换为数值类型
val labelIndexer = new StringIndexer()
  .setInputCol("label")
  .setOutputCol("indexedLabel")

模型训练

现在我们准备好开始构建和训练我们的LightGBM分类器了。我们将设定目标函数为多分类，并划分数据集为训练集和测试集。

// 创建LightGBM分类器，并设置为多分类
val lgbm = new LightGBMClassifier()
  .setLabelCol("indexedLabel")
  .setFeaturesCol("features")
  .setObjective("multiclass") // 设置目标函数为多分类

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

// 构建Pipeline
val pipeline = new Pipeline().setStages(Array(labelIndexer, assembler, lgbm))

// 训练模型
val model = pipeline.fit(trainingData)

模型评估

最后，我们在测试集上进行预测，并使用MulticlassClassificationEvaluator评估模型性能。

// 在测试集上进行预测
val predictions = model.transform(testData)

// 使用MulticlassClassificationEvaluator评估模型性能
val evaluator = new MulticlassClassificationEvaluator()
  .setLabelCol("indexedLabel")
  .setPredictionCol("prediction")
  .setMetricName("accuracy") // 可以选择其他的评价指标如"f1"

val accuracy = evaluator.evaluate(predictions)
println(s"The accuracy for test set is $accuracy")