[机器学习+pyspark库包]随机森林模型训练(实战)

import_random
129 阅读4分钟

我们将使用pyspark的MLlib库来构建一个随机森林分类模型。

步骤如下：

  1. 初始化SparkSession
  2. 生成模拟数据（特征和标签）
  3. 将数据转换为DataFrame
  4. 数据预处理：将特征向量化（使用VectorAssembler）
  5. 划分训练集和测试集
  6. 定义随机森林模型
  7. 训练模型
  8. 在测试集上评估模型效果
  9. 可选：模型参数调优（这里不展开，但可以提及）
  10. 可选：保存和加载模型注意：由于是模拟数据，我们生成一个二分类问题。

下面我们详细展开每一步，并给出代码。

下面是一个使用 PySpark MLlib 训练随机森林模型的详细步骤指南，包含模拟数据生成和完整代码：

步骤 1：环境准备

from pyspark.sql import SparkSession
from pyspark.ml.feature import VectorAssembler
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.evaluation import MulticlassClassificationEvaluator
import numpy as np
import pandas as pd

# 初始化 Spark 会话
spark = SparkSession.builder \
    .appName("RandomForestExample") \
    .config("spark.driver.memory", "4g") \  # 根据需求调整内存
    .getOrCreate()

步骤 2：生成模拟数据

# 生成 10000 个样本的模拟数据
np.random.seed(42)
n_samples = 10000

# 创建特征列 (f1-f5) 和标签列 (label)
data = {
    'f1': np.random.normal(0, 1, n_samples),
    'f2': np.random.poisson(5, n_samples),
    'f3': np.random.uniform(0, 10, n_samples),
    'f4': np.random.chisquare(3, n_samples),
    'f5': np.random.binomial(1, 0.3, n_samples),
    'label': np.random.binomial(1, 0.6, n_samples)  # 二元分类标签
}

# 创建 Pandas DataFrame 并转换为 Spark DataFrame
pdf = pd.DataFrame(data)
print(pdf)
df = spark.createDataFrame(pdf)
print(df)

# 查看数据结构
df.printSchema()
df.show(5)

步骤 3：数据预处理

# 合并特征列到单个特征向量
feature_cols = ['f1', 'f2', 'f3', 'f4', 'f5']
assembler = VectorAssembler(inputCols=feature_cols, outputCol="features")
df = assembler.transform(df)

# 选择需要的列
df = df.select("features", "label")

# 查看处理后的数据
df.show(5, truncate=False)

步骤 4：划分训练集/测试集

(train_data, test_data) = df.randomSplit([0.8, 0.2], seed=42)
print(f"训练集数量: {train_data.count()}")
print(f"测试集数量: {test_data.count()}")

步骤 5：配置并训练随机森林模型

# 初始化随机森林分类器
rf = RandomForestClassifier(
    labelCol="label",
    featuresCol="features",
    numTrees=100,               # 树的数量
    maxDepth=5,                 # 最大深度
    impurity="gini",            # 分裂准则 (gini/entropy)
    featureSubsetStrategy="auto",  # 特征选择策略
    seed=42,
    subsamplingRate=0.8         # 每棵树使用的数据比例
)

# 训练模型
model = rf.fit(train_data)

步骤 6：模型评估

# 在测试集上进行预测
predictions = model.transform(test_data)

# 显示预测结果
predictions.select("label", "prediction", "probability").show(10)

# 计算评估指标
evaluator = MulticlassClassificationEvaluator(
    labelCol="label",
    predictionCol="prediction",
    metricName="accuracy"
)

accuracy = evaluator.evaluate(predictions)
print(f"测试集准确率 = {accuracy:.4f}")

# 计算 AUC (需要二元分类)
from pyspark.ml.evaluation import BinaryClassificationEvaluator
auc_evaluator = BinaryClassificationEvaluator(
    labelCol="label", rawPredictionCol="rawPrediction")
auc = auc_evaluator.evaluate(predictions)
print(f"AUC = {auc:.4f}")

# 查看特征重要性
feature_importances = model.featureImportances.toArray()
for i, imp in enumerate(feature_importances):
    print(f"特征 f{i+1} 重要性: {imp:.4f}")

步骤 7：模型保存与加载（可选）

# 保存模型
model.save("hdfs:///user/models/rf_model")

# 加载模型
from pyspark.ml.classification import RandomForestClassificationModel
loaded_model = RandomForestClassificationModel.load("hdfs:///user/models/rf_model")

完整代码示例

from pyspark.sql import SparkSession
from pyspark.ml.feature import VectorAssembler
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.evaluation import MulticlassClassificationEvaluator, BinaryClassificationEvaluator
import numpy as np
import pandas as pd

# 1. 初始化 Spark
spark = SparkSession.builder \
    .appName("RandomForestDemo") \
    .config("spark.executor.memory", "2g") \
    .getOrCreate()

# 2. 生成模拟数据
np.random.seed(42)
data = {
    'f1': np.random.normal(0, 1, 10000),
    'f2': np.random.poisson(5, 10000),
    'f3': np.random.uniform(0, 10, 10000),
    'f4': np.random.chisquare(3, 10000),
    'f5': np.random.binomial(1, 0.3, 10000),
    'label': np.random.binomial(1, 0.6, 10000)
}
df = spark.createDataFrame(pd.DataFrame(data))

# 3. 特征工程
assembler = VectorAssembler(
    inputCols=['f1','f2','f3','f4','f5'],
    outputCol="features"
)
df = assembler.transform(df).select("features", "label")

# 4. 数据拆分
train, test = df.randomSplit([0.8, 0.2], seed=42)

# 5. 训练模型
rf = RandomForestClassifier(
    numTrees=150,
    maxDepth=7,
    featureSubsetStrategy="sqrt",
    impurity="entropy",
    subsamplingRate=0.75,
    seed=42
)
model = rf.fit(train)

# 6. 预测评估
predictions = model.transform(test)
accuracy = MulticlassClassificationEvaluator(metricName="accuracy").evaluate(predictions)
auc = BinaryClassificationEvaluator().evaluate(predictions)

print(f"\n{'='*50}")
print(f"模型评估结果:")
print(f"测试集准确率: {accuracy:.4f}")
print(f"AUC 值: {auc:.4f}")

# 7. 特征重要性
print("\n特征重要性:")
features = ['f1','f2','f3','f4','f5']
importances = model.featureImportances
for i, imp in enumerate(importances):
    print(f"{features[i]}: {imp:.4f}")

# 8. 保存模型
model.write().overwrite().save("random_forest_model")

spark.stop()

关键参数说明：

  1. numTrees：森林中树的数量（增加可提升性能但会增加计算成本）
  2. maxDepth：单棵树的最大深度（控制模型复杂度）
  3. impurity：分裂标准（gini 或 entropy）
  4. featureSubsetStrategy：特征采样策略（常用 "auto"、"sqrt"、"log2"）
  5. subsamplingRate：每棵树训练时的行采样比例

性能优化建议：

  1. 增加 numTrees（通常 100-500 效果较好）
  2. 使用交叉验证进行超参数调优
  3. 对于大数据集，增加 subsamplingRate
  4. 使用 cache() 缓存重复使用的 DataFrame
  5. 调整 Spark 资源配置（executor 内存/核心数）

典型输出示例：

模型评估结果:
测试集准确率: 0.7243
AUC 值: 0.7812

特征重要性:
f1: 0.2114
f2: 0.1857
f3: 0.1982
f4: 0.2059
f5: 0.1988

这个流程涵盖了从数据生成到模型评估的完整过程，您可以根据实际需求调整数据规模、特征数量和模型参数。

avatar
文章
阅读
粉丝