Scala下使用SparkML Pipeline实现训练/测试集标准化的方案咨询

方案1：单Pipeline全流程实现（更推荐） #

Spark Pipeline的设计原生就符合避免测试集数据泄露的规范，无需单独提取标准化模型：

// 数据集拆分
val training = splitData(0)
val test = splitData(1)

// 定义流程组件
val assembler = new VectorAssembler()
  .setInputCols(training.columns)
  .setOutputCol("features")   
val standardScaler = new StandardScaler()
  .setInputCol("features")
  .setOutputCol("normFeatures")
  .setWithStd(true)
  .setWithMean(true)
val lr = new LinearRegression().setFeaturesCol("normFeatures")

// 组装并拟合Pipeline，拟合阶段仅使用训练集计算所有转换器、模型参数
val pipeline = new Pipeline().setStages(Array(assembler, standardScaler, lr))
val pipelineModel = pipeline.fit(training)

// 直接对测试集做转换，内部自动复用训练阶段得到的标准化参数处理测试集，无偏差风险
val testPrediction = pipelineModel.transform(test)

如果确实有单独获取训练后标准化模型的需求，可以直接从拟合后的PipelineModel中提取：

val trainedScaler = pipelineModel.stages(1).asInstanceOf[StandardScalerModel]

方案2：拆分Pipeline的实现完全正确 #

你当前写的拆分预处理Pipeline和训练Pipeline的逻辑完全合规，既保证了标准化参数仅通过训练集计算，也满足你单独调用预处理模型转换数据的需求，可正常使用。

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