Pyspark k重交叉验证平均RMSE

from pyspark.sql import SparkSession
from pyspark.ml.regression import LinearRegression
from pyspark.ml.feature import VectorAssembler
from pyspark.ml.evaluation.RegressionEvaluator
from pyspark.ml.tuning.CrossValidator, CrossValidatorModel
from pyspark.ml.tuning.ParamGridBuilder

# 创建Spark会话
spark = SparkSession.builder.appName("K-Fold Cross Validation").getOrCreate()

# 加载数据
data = spark.read.csv("your_data.csv", header=True, inferSchema=True)

# 假设数据中包含特征列和一个目标列"target"
feature_cols = data.columns[:-1]
assembler = VectorAssembler(inputCols=feature_cols, outputCol="features")
data = assembler.transform(data).select("features", "target")

# 创建线性回归模型
lr = LinearRegression(featuresCol="features", labelCol="target")

# 创建评估器
evaluator = RegressionEvaluator(labelCol="target", predictionCol="prediction", metricName="rmse")

# 创建参数网格
paramGrid = ParamGridBuilder() \
    .addGrid(lr.regParam, [0.1, 0.01]) \
    .addGrid(lr.elasticNetParam, [0.0, 0.5, 1.0]) \
    .build()

# 创建k重交叉验证对象
k = 5
crossValidator = CrossValidator(estimator=lr,
                                  estimatorParamMaps=paramGrid,
                                  evaluator=evaluator,
                                  numFolds=k)

# 训练模型
cvModel = crossValidator.fit(data)

# 获取最佳模型
bestModel = cvModel.bestModel

# 使用最佳模型进行预测
predictions = bestModel.transform(data)

# 计算RMSE
rmse = evaluator.evaluate(predictions)
print("Root Mean Squared Error (RMSE):", rmse)

# 关闭Spark会️话
spark.stop()