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(_))
   }
 }

结果：