Keras小批量训练时如何为不平衡标签调整单样本梯度/学习率权重

你要的给标签1样本梯度提升9倍的需求无需修改fit底层源码，Keras提供了原生支持的实现方式，成本从低到高有以下几种可选：

1. 直接使用fit方法的class_weight参数（最推荐） #

这个是fit方法原生支持的参数，专门用于解决类别不平衡场景的损失加权问题，刚好匹配你的需求，仅需要在你原有代码的基础上新增一行参数即可：

history = model.fit(
    trainX, trainY, 
    epochs=1, 
    batch_size=minibatch_size, 
    validation_data=(valX, valY), 
    verbose=0,
    # 标签0权重为1，标签1权重为9，损失计算时会自动对对应样本的损失乘权重
    class_weight={0: 1.0, 1: 9.0}
)

该方案的原理是计算损失时给标签1的样本损失放大9倍，反向传播时对应梯度也会同步放大9倍，和你要求的给标签1样本设置9倍学习率的效果完全等价，不需要改动任何其他逻辑。

2. 自定义加权损失函数 #

如果你需要更灵活的加权规则，可以自己实现加权的损失函数，编译模型时传入即可，以二元交叉熵为例：

import tensorflow as tf

def weighted_binary_crossentropy(y_true, y_pred):
    base_bce = tf.keras.losses.binary_crossentropy(y_true, y_pred)
    # 标签1的损失乘9，标签0保持不变
    weight = tf.where(y_true == 1, 9.0, 1.0)
    return tf.reduce_mean(base_bce * weight)

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

后续你原来的fit调用代码不需要做任何修改即可生效。

3. 重写train_step自定义训练逻辑（对应你说的修改fit底层逻辑的需求） #

如果需要完全控制训练过程中的梯度计算规则，可以继承Keras的Model类重写train_step方法，自定义梯度加权逻辑：

import tensorflow as tf

class GradientWeightedModel(tf.keras.Model):
    def train_step(self, data):
        x, y = data
        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)
            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
            # 自定义梯度加权
            weight_mask = tf.where(y == 1, 9.0, 1.0)
            loss = tf.reduce_mean(loss * weight_mask)
        # 计算并应用梯度
        gradients = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
        # 更新训练指标
        self.compiled_metrics.update_state(y, y_pred)
        return {m.name: m.result() for m in self.metrics}

构建模型时使用上述自定义的GradientWeightedModel类即可，后续fit调用逻辑不需要修改。

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