PyTorch多分类任务用sklearn compute_class_weight计算类权重该用训练集还是全量数据

核心结论 #

计算类权重时仅能使用训练集数据，不可引入验证集、测试集的任何数据信息。 #

原因说明 #

• 类权重是训练阶段用来缓解类别不均衡带来的学习偏差的优化手段，属于模型训练流程的一部分，必须严格遵守「训练过程不得访问评估阶段数据」的基础规范，一旦用到验证/测试集的类别分布信息就会造成数据泄露，导致最终的评估指标失真，无法正确反映模型的真实泛化能力。
• 验证集、测试集的核心作用是模拟模型上线后面对的未知数据，提前获取它们的分布信息相当于在评估前给模型“透题”，完全违背了数据集划分的初衷。

PyTorch结合sklearn的实操示例 #

仅传入训练集标签计算权重即可：

from sklearn.utils.class_weight import compute_class_weight
import numpy as np
import torch

# y_train为训练集的标签数组
class_weights = compute_class_weight(
    class_weight="balanced",
    classes=np.unique(y_train),
    y=y_train
)
# 转换为PyTorch损失函数可用的张量格式
class_weights = torch.tensor(class_weights, dtype=torch.float32)

# 传入交叉熵损失函数即可使用
loss_fn = torch.nn.CrossEntropyLoss(weight=class_weights)

补充说明 #

如果你是通过分层抽样的方式划分数据集，训练集、验证集、测试集的类别分布完全一致，此时用全量数据计算权重和仅用训练集计算的结果差异极小，但依然不建议这么做，严格遵守训练与评估数据隔离的规范能避免额外的实验误差。

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