如何在Apache Spark MLlib(Python)中为逻辑回归模型分配类权重

结论 #

你描述的是Spark MLlib处理类别不平衡问题的标准正确方案，权重赋值逻辑完全合理。

原理解释 #

Spark MLlib中LogisticRegression的weightCol参数作用是为每一行样本设置训练时的损失权重：权重越高的样本，模型在迭代优化过程中越会优先拟合其分类结果，避免默认情况下模型向占比更高的类别倾斜。
针对你的场景：负类（0.0）占比75%，正类（1.0）占比25%，你将正类权重设为0.75、负类权重设为0.25，刚好让两个类别的总样本权重之和相等，模型会平等学习两个类别的特征，不会出现绝大多数样本都被预测为负类的偏差问题。

注意：权重的绝对值不影响最终训练效果，只要不同类别之间的权重比例符合类别占比的反比即可，你也可以设置正类权重为3、负类权重为1，最终训练效果和当前赋值一致。

参考实现代码 #

你可以直接参考以下PySpark代码实现：

from pyspark.sql.functions import when
from pyspark.ml.classification import LogisticRegression

# 为原始数据集添加权重列
df_with_weight = df.withColumn(
    "weight",
    when(df.label == 1.0, 0.75).otherwise(0.25)
)

# 初始化逻辑回归模型，指定权重列
blor = LogisticRegression(weightCol="weight", labelCol="label")

# 后续按正常流程训练模型即可
lr_model = blor.fit(df_with_weight)

拓展优化建议 #

如果你的数据集类别占比后续可能发生变化，也可以自动统计类别占比生成权重，避免手动写死数值：

# 统计各类别样本量
class_count = df.groupBy("label").count().collect()
total_samples = df.count()
# 按占比反比生成权重映射
weight_map = {row["label"]: (total_samples - row["count"]) / total_samples for row in class_count}

# 批量生成权重列
df_with_weight = df.replace(weight_map, subset=["label"]).withColumnRenamed("label", "weight")
# 再合并回原始标签列即可

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