[Case Study]Decision Tree Classification using Scala

是 否 +2 论坛币

k人 参与回答

经管之家送您一份

应届毕业生专属福利!

求职就业群

赵安豆老师微信：zhaoandou666

经管之家联合CDA

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

立即领取

感谢您参与论坛问题回答

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

+2 论坛币

1. package org.apache.spark.examples.ml
   
2. import org.apache.spark.sql.SQLContext
   
3. import org.apache.spark.{SparkContext, SparkConf}
   
4. // $example on$
   
5. import org.apache.spark.ml.Pipeline
   
6. import org.apache.spark.ml.classification.DecisionTreeClassifier
   
7. import org.apache.spark.ml.classification.DecisionTreeClassificationModel
   
8. import org.apache.spark.ml.feature.{StringIndexer, IndexToString, VectorIndexer}
   
9. import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
   
10. // $example off$
    
11. 
    
12. object DecisionTreeClassificationExample {
    
13.   def main(args: Array[String]): Unit = {
    
14.     val conf = new SparkConf().setAppName("DecisionTreeClassificationExample")
    
15.     val sc = new SparkContext(conf)
    
16.     val sqlContext = new SQLContext(sc)
    
17.     // $example on$
    
18.     // Load the data stored in LIBSVM format as a DataFrame.
    
19.     val data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
    
20. 
    
21.     // Index labels, adding metadata to the label column.
    
22.     // Fit on whole dataset to include all labels in index.
    
23.     val labelIndexer = new StringIndexer()
    
24.       .setInputCol("label")
    
25.       .setOutputCol("indexedLabel")
    
26.       .fit(data)
    
27.     // Automatically identify categorical features, and index them.
    
28.     val featureIndexer = new VectorIndexer()
    
29.       .setInputCol("features")
    
30.       .setOutputCol("indexedFeatures")
    
31.       .setMaxCategories(4) // features with > 4 distinct values are treated as continuous
    
32.       .fit(data)
    
33. 
    
34.     // Split the data into training and test sets (30% held out for testing)
    
35.     val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
    
36. 
    
37.     // Train a DecisionTree model.
    
38.     val dt = new DecisionTreeClassifier()
    
39.       .setLabelCol("indexedLabel")
    
40.       .setFeaturesCol("indexedFeatures")
    
41. 
    
42.     // Convert indexed labels back to original labels.
    
43.     val labelConverter = new IndexToString()
    
44.       .setInputCol("prediction")
    
45.       .setOutputCol("predictedLabel")
    
46.       .setLabels(labelIndexer.labels)
    
47. 
    
48.     // Chain indexers and tree in a Pipeline
    
49.     val pipeline = new Pipeline()
    
50.       .setStages(Array(labelIndexer, featureIndexer, dt, labelConverter))
    
51. 
    
52.     // Train model.  This also runs the indexers.
    
53.     val model = pipeline.fit(trainingData)
    
54. 
    
55.     // Make predictions.
    
56.     val predictions = model.transform(testData)
    
57. 
    
58.     // Select example rows to display.
    
59.     predictions.select("predictedLabel", "label", "features").show(5)
    
60. 
    
61.     // Select (prediction, true label) and compute test error
    
62.     val evaluator = new MulticlassClassificationEvaluator()
    
63.       .setLabelCol("indexedLabel")
    
64.       .setPredictionCol("prediction")
    
65.       .setMetricName("precision")
    
66.     val accuracy = evaluator.evaluate(predictions)
    
67.     println("Test Error = " + (1.0 - accuracy))
    
68. 
    
69.     val treeModel = model.stages(2).asInstanceOf[DecisionTreeClassificationModel]
    
70.     println("Learned classification tree model:\n" + treeModel.toDebugString)
    
71.     // $example off$
    
72.   }
    
73. }

复制代码

扫码加我 拉你入群

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

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

分享0 收藏0 回帖

关键词：Case study Decision cation Using SCALA example package import