基于自定义损失函数处理多分类任务中的类别不平衡问题

我正在用Python结合先进机器学习技术解决多分类任务，当前数据集存在严重类别不平衡问题，少数类样本占比极低，导致模型对少数类的性能表现不佳。我采用TensorFlow/Keras开发模型，当前模型结构如下：

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

# Example model architecture
model = Sequential([
    Dense(128, activation='relu', input_shape=(input_shape,)),
    Dense(64, activation='relu'),
    Dense(num_classes, activation='softmax')
])

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

其中num_classes代表数据集类别数，input_shape为输入特征维度。现有损失函数categorical_crossentropy无法适配类别不平衡问题，因此我希望构建融入类别权重的自定义损失函数，提升少数类的训练权重。具体问题如下：

1. 如何在TensorFlow/Keras中开发适用于多分类任务、融入类别权重的自定义损失函数？
2. 有哪些最佳实践可确保该自定义损失函数计算高效，不会显著增加训练时间？
3. 实现此类自定义损失函数处理类别不平衡时，有哪些潜在陷阱或常见错误需要注意？

1. 自定义带类别权重的多分类损失函数实现 #

首先需要根据数据集的类别分布计算类别权重，常用的计算方式有两种：

• 基于样本数的倒数：class_weight = total_samples / (num_classes * class_samples)
• 有效样本数权重（避免极端值）：class_weight = tf.math.log(total_samples / (class_samples + 1))

针对one-hot编码标签（对应categorical_crossentropy），可以按以下方式实现自定义损失：

import tensorflow as tf

# 第一步：计算类别权重（假设已统计每个类别的样本数class_samples）
total_samples = sum(class_samples)
class_weights = total_samples / (len(class_samples) * class_samples)
# 转为TensorFlow常量，避免重复计算
class_weights = tf.convert_to_tensor(class_weights, dtype=tf.float32)

# 第二步：定义自定义损失函数
def weighted_categorical_crossentropy(y_true, y_pred):
    # 限制y_pred范围，避免log(0)的数值问题
    y_pred_clipped = tf.clip_by_value(y_pred, 1e-7, 1.0)
    # 计算每个样本的交叉熵损失
    cross_entropy = -y_true * tf.math.log(y_pred_clipped)
    # 按类别加权
    weighted_loss = cross_entropy * class_weights
    # 求平均损失
    return tf.reduce_mean(tf.reduce_sum(weighted_loss, axis=-1))

# 编译模型时使用自定义损失
model.compile(optimizer='adam', loss=weighted_categorical_crossentropy, metrics=['accuracy'])

如果使用稀疏整数标签（对应sparse_categorical_crossentropy），可以通过tf.gather匹配对应类别的权重：

def weighted_sparse_categorical_crossentropy(y_true, y_pred):
    # 获取每个样本对应类别的权重
    sample_weights = tf.gather(class_weights, tf.cast(y_true, tf.int32))
    # 使用内置交叉熵函数，保证数值稳定性
    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y_true, logits=y_pred)
    # 加权后求平均
    return tf.reduce_mean(cross_entropy * sample_weights)

2. 确保损失函数计算高效的最佳实践 #

• 提前固化权重：将类别权重转为TensorFlow常量或非可训练变量，避免在损失函数内部重复计算统计量，减少额外开销。
• 优先使用内置交叉熵函数：tf.nn.softmax_cross_entropy_with_logits_v2或tf.nn.sparse_softmax_cross_entropy_with_logits经过官方优化，数值稳定性更好，且能充分利用GPU并行计算，比手动实现交叉熵效率更高。
• 全向量化操作：避免使用Python循环处理样本或类别，全部用TensorFlow张量操作实现，最大化并行计算能力。
• 封装权重为闭包参数：把预计算好的权重作为外部参数传入损失函数，不要在函数内部执行权重计算逻辑，简化函数内部流程。
• 验证集使用原始损失：训练阶段用加权损失，验证/测试阶段切换回普通交叉熵，避免不必要的加权计算，同时保证评估指标的真实性。

3. 实现时的潜在陷阱与常见错误 #

• 权重计算失衡：直接使用样本数倒数可能导致少数类权重过大，模型过度拟合少数类；建议采用有效样本数权重，或对权重做归一化处理（如除以权重最大值）。
• 标签格式混用：one-hot标签和稀疏标签的权重匹配逻辑完全不同，若混用会导致加权失效，比如用稀疏标签的gather逻辑处理one-hot标签，会出现权重对应错误。
• 数值稳定性问题：手动计算交叉熵时，若不对y_pred做范围限制，会出现log(0)导致的无穷大值，必须用tf.clip_by_value处理，或直接使用内置交叉熵函数。
• 过度加权少数类：少数类权重过高会让模型完全偏向少数类，牺牲多数类的性能；建议通过验证集调整权重系数，比如在权重前乘以0.5~1之间的系数平衡影响。
• 验证集误用加权损失：验证阶段使用加权损失会导致评估指标失真，无法反映模型在真实数据分布下的性能，必须用原始交叉熵计算验证损失。
• 混入非TensorFlow操作：在损失函数中使用Python列表、循环等非张量操作，会破坏自动微分流程，导致训练报错或大幅降低效率。

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