Keras 自定义 train_step 怎么让正则损失继续参与梯度

自定义 Keras train_step 后，正则损失不一定会自动进入你手写的梯度计算。关键做法是：前向调用模型后读取 self.losses，把它和 self.compute_loss() 返回的任务损失相加，再对总损失求梯度。这样 kernel_regularizer、activity_regularizer 或层里的 add_loss() 才会继续影响参数更新。

官方地址：https://keras.io/

症状：训练能跑，正则项却没有参与更新

常见场景是把官方的自定义训练步骤改成自己的损失函数：前向得到 y_pred，再计算 MSE 或交叉熵，随后直接执行 tape.gradient(task_loss, trainable_vars)。这段代码本身没有语法问题，但它只优化任务损失；模型前向期间由正则器或 add_loss() 注册的项被遗漏了。

Keras 的权重正则器在层调用后生成损失，官方 API 也明确说明应在前向之后访问 layer.losses。在模型层级的自定义训练步骤里，对应的集合就是 self.losses。因此排查顺序应先看前向位置，再看总损失是否在梯度带内完成。

根因：任务损失和 model.losses 是两条来源

可以把一批数据的目标写成下面的关系：

total_loss = task_loss(y, y_pred) + sum(model.losses)

其中 task_loss 通常来自 compile(loss=...)，而 model.losses 来自本次前向传播产生的正则项。它们不是跨 batch 累加的历史列表；顶层模型每次新的前向调用都会刷新这批附加损失。正因为如此，读取位置必须紧跟 self(x, training=True)，而不能提前缓存。

修复动作：在 GradientTape 内合并正则损失

下面的实现保留 fit()、回调和已配置指标，只改写每批训练的损失组成。代码中的 tf.add_n 需要至少一个张量，所以对空列表使用零张量；当模型没有正则项时，代码仍能沿用相同路径。

import tensorflow as tf
import keras

class RegularizedModel(keras.Model):
    def train_step(self, data):
        # data 可能是 (x, y)，需要权重时也可能是 (x, y, sample_weight)。
        if len(data) == 3:
            x, y, sample_weight = data
        else:
            x, y = data
            sample_weight = None

        with tf.GradientTape() as tape:
            # 先前向调用，self.losses 才会得到本批正则项。
            y_pred = self(x, training=True)
            task_loss = self.compute_loss(
                x=x, y=y, y_pred=y_pred, sample_weight=sample_weight
            )
            # 空列表时补零，避免 tf.add_n([]) 报错。
            regularization_loss = (
                tf.add_n(self.losses) if self.losses else tf.constant(0.0)
            )
            total_loss = task_loss + regularization_loss

        # 必须对包含正则项的 total_loss 求梯度。
        gradients = tape.gradient(total_loss, self.trainable_variables)
        self.optimizer.apply_gradients(
            zip(gradients, self.trainable_variables)
        )

        # 日志中的 loss 与真实优化目标保持一致；其他指标照常更新。
        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(total_loss)
            else:
                metric.update_state(y, y_pred, sample_weight=sample_weight)
        return {metric.name: metric.result() for metric in self.metrics}

如果项目不使用 sample_weight，可以删掉解包和传参部分，但不要删掉 self.losses 的合并。对于 Keras 3 多后端模型，tf.GradientTape 这段写法属于 TensorFlow 后端实现；若要切换 JAX 或 PyTorch，应按对应后端重写训练步骤，不能只替换导入。

防复发：把梯度、指标和边界一起对齐

修复后重点检查三件事。第一，正则器必须挂在实际被调用的层上，self.losses 要在本次前向之后读取。第二，梯度、日志和回调看到的是同一个 total_loss；只把正则项加进梯度、却继续记录 task_loss，会让排障误判。第三，activity_regularizer 的数值会按输入 batch size 处理，不能把它再按样本数重复缩放。

常见问题

为什么直接使用 self.compute_loss 还不够？

它主要负责 compile() 配置的任务损失；正则器或 add_loss() 生成的附加项仍需要从 self.losses 取出并合并。

self.losses 会不会把上一批的正则项带进来？

正常前向调用会刷新顶层模型的附加损失，集合对应最近一次前向。不要在一次前向后延迟到下一批才读取，也不要把它当作训练历史。

为什么日志里的 loss 比预期大？

日志若记录的是 total_loss，它本来就包含任务损失和正则项。应分别记录调试用的正则项与总损失，而不是为了让数字变小而移除正则项。