PySpark分批训练朴素贝叶斯模型：如何避免模型被覆盖？

首先得明确：Spark ML中的NaiveBayes并没有像Scikit-learn那样提供partial_fit方法，所以你每次调用pipeline.fit()都会从头训练一个全新的模型，自然会覆盖之前的结果。不过因为多项式NaiveBayes是基于统计计数的模型，我们可以手动累加各批次的统计量，来实现增量训练的效果。下面是具体的解决方案：

步骤1：拆分预处理管道与模型训练 #

你的Pipeline里包含了预处理（StringIndexer、Tokenizer等）和模型训练（NaiveBayes）两个部分。预处理阶段的参数（比如停用词列表、哈希特征维度）是固定的，而且StringIndexer的标签映射最好保持一致，所以我们可以先单独拟合预处理管道，后续所有批次都复用这个预处理模型。

from pyspark.ml import Pipeline
from pyspark.ml.feature import StringIndexer, Tokenizer, StopWordsRemover, HashingTF

# 定义预处理阶段
categoryIndexer = StringIndexer(inputCol="diff", outputCol="label")
tokenizer = Tokenizer(inputCol="text", outputCol="raw")
remover = StopWordsRemover(inputCol="raw", outputCol="words")
hashingTF = HashingTF(inputCol="words", outputCol="features", numFeatures=100000)

preprocessing_stages = [categoryIndexer, tokenizer, remover, hashingTF]
preprocessing_pipeline = Pipeline(stages=preprocessing_stages)

# 用第一个批次的训练数据拟合预处理管道（确保标签映射固定）
first_df = sqlContext.read.jdbc(url=url,table="(SELECT text, diff FROM tweets INNER JOIN djitf ON tweets.created = djitf.day WHERE id > 0 AND id < 1000000) as table1", properties=properties)
first_train, _ = first_df.randomSplit([0.8, 0.2])
preprocessing_model = preprocessing_pipeline.fit(first_train)

步骤2：手动累加NaiveBayes的核心统计量 #

多项式NaiveBayes的核心依赖两个统计量：

1. 每个类别的样本数量（用于计算先验概率）
2. 每个类别下所有特征的总出现次数（用于计算条件概率）

我们可以维护全局的统计变量，在每个批次中更新这些变量：

import math
from pyspark.ml.linalg import Vectors
from pyspark.sql import functions as F

# 初始化全局统计量：先从第一个批次获取所有类别
first_transformed = preprocessing_model.transform(first_train)
categories = first_transformed.select("label").distinct().rdd.map(lambda x: x[0]).collect()

# 初始化类别计数和特征总和
class_counts = {cat: 0 for cat in categories}
feature_sums = {cat: Vectors.zeros(100000) for cat in categories}  # 与HashingTF的numFeatures一致
smoothing = 1.0  # 和你原来的NaiveBayes参数保持一致

# 遍历所有批次
for i in range(0, 365):
    # 读取当前批次数据
    df = sqlContext.read.jdbc(
        url=url,
        table=f"(SELECT text, diff FROM tweets INNER JOIN djitf ON tweets.created = djitf.day WHERE id > {i*1000000} AND id < {(i+1)*1000000}) as table1",
        properties=properties
    )
    train_data, test_data = df.randomSplit([0.8, 0.2])
    
    # 用预处理模型转换数据
    transformed_train = preprocessing_model.transform(train_data)
    
    # 计算当前批次的统计量
    batch_stats = transformed_train.groupBy("label").agg(
        F.count("label").alias("count"),
        F.sum("features").alias("sum_features")
    ).collect()
    
    # 累加至全局统计量
    for row in batch_stats:
        cat = row["label"]
        class_counts[cat] += row["count"]
        feature_sums[cat] += row["sum_features"]

步骤3：用累加的统计量构建最终模型 #

当所有批次的统计量都累加完成后，我们可以手动计算NaiveBayes的logPrior和logProb参数，然后构建正式的NaiveBayesModel：

from pyspark.ml.classification import NaiveBayesModel

total_samples = sum(class_counts.values())
num_features = 100000

logPrior = []
logProb = []

# 计算每个类别的log先验和log条件概率
for cat in categories:
    # 计算logPrior（加入平滑项）
    prior = (class_counts[cat] + smoothing) / (total_samples + smoothing * len(categories))
    logPrior.append(math.log(prior))
    
    # 计算该类别下的总特征数（加入平滑项）
    total_feature = feature_sums[cat].sum() + smoothing * num_features
    
    # 计算每个特征的log概率
    prob_vec = feature_sums[cat].map(lambda x: math.log((x + smoothing) / total_feature))
    logProb.append(prob_vec)

# 构建最终的增量训练模型
final_nb_model = NaiveBayesModel(
    labelCol="label",
    featuresCol="features",
    logPrior=logPrior,
    logProb=logProb,
    modelType="multinomial"
)

关键注意事项 #

• 标签一致性：确保所有批次的diff列类别都在初始的类别列表中，如果有新类别出现，需要提前处理（比如给StringIndexer设置handleInvalid="keep"，或者先全量查询所有类别）。
• 特征维度固定：HashingTF的numFeatures必须全程保持一致，否则特征向量维度不匹配，无法累加。
• 平滑参数统一：平滑参数smoothing要和你最初的NaiveBayes设置一致，保证概率计算的一致性。
• 仅适用于多项式NB：这种手动累加的方法只适合多项式NaiveBayes，高斯NB等其他类型不适用，因为它们的统计逻辑不同。

内容的提问来源于stack exchange，提问作者Torphexi