#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""
Decision tree classification and regression using MLlib.
This example requires NumPy (http://www.numpy.org/).
"""
import numpy, os, sys
from operator import add
from pyspark import SparkContext
from pyspark.mllib.regression import LabeledPoint
from pyspark.mllib.tree import DecisionTree
from pyspark.mllib.util import MLUtils
def getAccuracy(dtModel, data):
"""
Return accuracy of DecisionTreeModel on the given RDD[LabeledPoint].
"""
seqOp = (lambda acc, x: acc + (x[0] == x[1]))
predictions = dtModel.predict(data.map(lambda x: x.features))
truth = data.map(lambda p: p.label)
trainCorrect = predictions.zip(truth).aggregate(0, seqOp, add)
if data.count() == 0:
return 0
return trainCorrect / (0.0 + data.count())
def getMSE(dtModel, data):
"""
Return mean squared error (MSE) of DecisionTreeModel on the given
RDD[LabeledPoint].
"""
seqOp = (lambda acc, x: acc + numpy.square(x[0] - x[1]))
predictions = dtModel.predict(data.map(lambda x: x.features))
truth = data.map(lambda p: p.label)
trainMSE = predictions.zip(truth).aggregate(0, seqOp, add)
if data.count() == 0:
return 0
return trainMSE / (0.0 + data.count())
def reindexClassLabels(data):
"""
Re-index class labels in a dataset to the range {0,...,numClasses-1}.
If all labels in that range already appear at least once,
then the returned RDD is the same one (without a mapping).
Note: If a label simply does not appear in the data,
the index will not include it.
Be aware of this when reindexing subsampled data.
:param data: RDD of LabeledPoint where labels are integer values
denoting labels for a classification problem.
:return: Pair (reindexedData, origToNewLabels) where
reindexedData is an RDD of LabeledPoint with labels in
the range {0,...,numClasses-1}, and
origToNewLabels is a dictionary mapping original labels
to new labels.
"""
# classCounts: class --> # examples in class
classCounts = data.map(lambda x: x.label).countByValue()
numExamples = sum(classCounts.values())
sortedClasses = sorted(classCounts.keys())
numClasses = len(classCounts)
# origToNewLabels: class --> index in 0,...,numClasses-1
if (numClasses < 2):
print >> sys.stderr, \
"Dataset for classification should have at least 2 classes." + \
" The given dataset had only %d classes." % numClasses
exit(1)
origToNewLabels = dict([(sortedClasses[i], i) for i in range(0, numClasses)])
print "numClasses = %d" % numClasses
print "Per-class example fractions, counts:"
print "Class\tFrac\tCount"
for c in sortedClasses:
frac = classCounts[c] / (numExamples + 0.0)
print "%g\t%g\t%d" % (c, frac, classCounts[c])
if (sortedClasses[0] == 0 and sortedClasses[-1] == numClasses - 1):
return (data, origToNewLabels)
else:
reindexedData = \
data.map(lambda x: LabeledPoint(origToNewLabels[x.label], x.features))
return (reindexedData, origToNewLabels)
def usage():
print >> sys.stderr, \
"Usage: decision_tree_runner [libsvm format data filepath]\n" + \
" Note: This only supports binary classification."
exit(1)
if __name__ == "__main__":
if len(sys.argv) > 2:
usage()
sc = SparkContext(appName="PythonDT")
# Load data.
dataPath = 'data/mllib/sample_libsvm_data.txt'
if len(sys.argv) == 2:
dataPath = sys.argv[1]
if not os.path.isfile(dataPath):
sc.stop()
usage()
points = MLUtils.loadLibSVMFile(sc, dataPath)
# Re-index class labels if needed.
(reindexedData, origToNewLabels) = reindexClassLabels(points)
# Train a classifier.
categoricalFeaturesInfo={} # no categorical features
model = DecisionTree.trainClassifier(reindexedData, numClasses=2,
categoricalFeaturesInfo=categoricalFeaturesInfo)
# Print learned tree and stats.
print "Trained DecisionTree for classification:"
print " Model numNodes: %d\n" % model.numNodes()
print " Model depth: %d\n" % model.depth()
print " Training accuracy: %g\n" % getAccuracy(model, reindexedData)
print model
sc.stop()