Spark线性回归实现优化

import org.apache.log4j.{Level, Logger}
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.ml.regression.LinearRegression
import org.apache.spark.sql.SparkSession

/**
  * 线性回归
  * Created by zhen on 2018/11/12.
  */
object LinearRegression {
  Logger.getLogger("org").setLevel(Level.WARN) // 设置日志级别
  def main(args: Array[String]) {
    val spark = SparkSession
      .builder()
      .appName("LinearRegression")
      .master("local[2]")
      .getOrCreate()
    val train_data = spark.sparkContext.textFile("E:/BDS/newsparkml/src/train.txt") // 加载数据
    val train_map_data = train_data.map{ row =>
        val split = row.split(",")
        (split(0).toDouble,split(1).toDouble,split(2).toDouble,split(3).toDouble,
          split(4).toDouble,split(5).toDouble,split(6).toDouble,split(7).toDouble)
      }
    val df = spark.sqlContext.createDataFrame(train_map_data)
    val colArray = Array("Population","Income","Illiteracy","LifeExp","HSGrad","Frost","Area")
    val train_df = df.toDF(colArray(0),colArray(1),colArray(2),colArray(3),"Murder",colArray(4),colArray(5),colArray(6))
    val assembler = new VectorAssembler()
      .setInputCols(colArray)
      .setOutputCol("features")
    val vectDF = assembler.transform(train_df)
    val weights = Array(0.8,0.2) //设置训练集和测试集的比例
    val split_data = vectDF.randomSplit(weights) // 拆分训练集和测试集
    // 创建模型对象
    val linearRegression = new LinearRegression()
      .setFeaturesCol("features")
      .setLabelCol("Murder")
      .setFitIntercept(true)
      .setMaxIter(10)
      .setRegParam(0.3)// 正则化
      .setElasticNetParam(0.8)
    // 训练模型
    val lrModel = linearRegression.fit(split_data(0))
    // 查看模型参数
    //lrModel.extractParamMap()
    println(s"Cofficients:${lrModel.coefficients} Intercept:${lrModel.intercept}")
    //模型评估
    val trainingSummary = lrModel.summary
    println(s"objectiveHistoryList:${trainingSummary.objectiveHistory.toList}")
    println(s"r2:${trainingSummary.r2}")
    // 预测
    val predictions = lrModel.transform(split_data(1))
    val predict_result = predictions.selectExpr("features","Murder","round(prediction,1) as prediction") // 保存一位小数
    println("训练集数据------------------------------真实值--预测值")
    predict_result.foreach(println(_))
  }
}

结果：