利用决策树算法进行预测

利用决策树算法进行预测

from pyspark.ml.feature import StringIndexer, VectorAssembler

df_train = spark.read.csv('./data/titanic-train.csv',header=True,inferSchema=True).cache()

df_train = df_train.fillna({'Age': round(29.699,0)})
df_train = df_train.fillna({'Embarked': 'S'})

labelIndexer = StringIndexer(inputCol="Embarked", outputCol="iEmbarked")
model = labelIndexer.fit(df_train)
df_train = model.transform(df_train)

labelIndexer = StringIndexer(inputCol="Sex", outputCol="iSex")
model = labelIndexer.fit(df_train)
df_train = model.transform(df_train)

# 选取数值类型的字段
features = ['Pclass', 'iSex', 'Age', 'SibSp', 'Parch', 'Fare', 'iEmbarked','Survived']
train_features = df_train[features]

# 将多列转换成向量
df_assembler = VectorAssembler(inputCols=['Pclass', 'iSex', 'Age', 'SibSp','Parch', 'Fare', 'iEmbarked'], outputCol="features")
train = df_assembler.transform(train_features)

from pyspark.ml.classification import DecisionTreeClassifier
#DecisionTree模型
dtree = DecisionTreeClassifier(labelCol="Survived", featuresCol="features")
treeModel = dtree.fit(train)

#打印treeModel
print(treeModel.toDebugString)

#利用模型的transform方法对训练数据train进行预测
dt_predictions = treeModel.transform(train)
# 从预测集合中选择 "prediction", "Survived", "features" 这三个字段进行显示
dt_predictions.select("prediction", "Survived", "features").show()

from pyspark.ml.evaluation import MulticlassClassificationEvaluator

multi_evaluator = MulticlassClassificationEvaluator(labelCol = 'Survived', metricName = 'accuracy')

print('Decision Tree Accu:', multi_evaluator.evaluate(dt_predictions))

注释：
accuracy 精度