[Scala Study]Decision Tree Regression Using Scala

是 否 +2 论坛币

k人 参与回答

经管之家送您一份

应届毕业生专属福利!

求职就业群

赵安豆老师微信：zhaoandou666

经管之家联合CDA

送您一个全额奖学金名额~ !

立即领取

感谢您参与论坛问题回答

经管之家送您两个论坛币！

+2 论坛币

1. // scalastyle:off println
   
2. package org.apache.spark.examples.ml
   
3. import org.apache.spark.sql.SQLContext
   
4. import org.apache.spark.{SparkContext, SparkConf}
   
5. // $example on$
   
6. import org.apache.spark.ml.Pipeline
   
7. import org.apache.spark.ml.regression.DecisionTreeRegressor
   
8. import org.apache.spark.ml.regression.DecisionTreeRegressionModel
   
9. import org.apache.spark.ml.feature.VectorIndexer
   
10. import org.apache.spark.ml.evaluation.RegressionEvaluator
    
11. // $example off$
    
12. object DecisionTreeRegressionExample {
    
13.   def main(args: Array[String]): Unit = {
    
14.     val conf = new SparkConf().setAppName("DecisionTreeRegressionExample")
    
15.     val sc = new SparkContext(conf)
    
16.     val sqlContext = new SQLContext(sc)
    
17. 
    
18.     // $example on$
    
19.     // Load the data stored in LIBSVM format as a DataFrame.
    
20.     val data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
    
21. 
    
22.     // Automatically identify categorical features, and index them.
    
23.     // Here, we treat features with > 4 distinct values as continuous.
    
24.     val featureIndexer = new VectorIndexer()
    
25.       .setInputCol("features")
    
26.       .setOutputCol("indexedFeatures")
    
27.       .setMaxCategories(4)
    
28.       .fit(data)
    
29. 
    
30.     // Split the data into training and test sets (30% held out for testing)
    
31.     val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
    
32. 
    
33.     // Train a DecisionTree model.
    
34.     val dt = new DecisionTreeRegressor()
    
35.       .setLabelCol("label")
    
36.       .setFeaturesCol("indexedFeatures")
    
37. 
    
38.     // Chain indexer and tree in a Pipeline
    
39.     val pipeline = new Pipeline()
    
40.       .setStages(Array(featureIndexer, dt))
    
41. 
    
42.     // Train model.  This also runs the indexer.
    
43.     val model = pipeline.fit(trainingData)
    
44. 
    
45.     // Make predictions.
    
46.     val predictions = model.transform(testData)
    
47. 
    
48.     // Select example rows to display.
    
49.     predictions.select("prediction", "label", "features").show(5)
    
50. 
    
51.     // Select (prediction, true label) and compute test error
    
52.     val evaluator = new RegressionEvaluator()
    
53.       .setLabelCol("label")
    
54.       .setPredictionCol("prediction")
    
55.       .setMetricName("rmse")
    
56.     val rmse = evaluator.evaluate(predictions)
    
57.     println("Root Mean Squared Error (RMSE) on test data = " + rmse)
    
58. 
    
59.     val treeModel = model.stages(1).asInstanceOf[DecisionTreeRegressionModel]
    
60.     println("Learned regression tree model:\n" + treeModel.toDebugString)
    
61.     // $example off$
    
62.   }
    
63. }

复制代码

扫码加我 拉你入群

请注明：姓名-公司-职位

以便审核进群资格，未注明则拒绝

关键词：regression regressio Decision regress SCALA example import

Scala for Data Science