TensorFlow自定义train_step结合多损失与正则化损失的正确性验证

你的代码完全符合需求，包括自定义train_step实现双损失线性组合、保留模型层的正则化损失，你对L2正则化的理解也是正确的——直接将带正则化的交叉熵损失与加权后的额外损失相加是合理且正确的做法。

核心逻辑验证 #

1. 正则化损失的正确整合

   • self.losses会自动收集模型所有层定义的正则化惩罚（比如Dense(kernel_regularizer=tf.keras.regularizers.L2())）。
   • 调用self.compiled_loss(y, y_pred, regularization_losses=self.losses)时，TensorFlow会自动将交叉熵损失与所有正则化损失求和，得到的ce_loss已经是包含L2惩罚的完整监督损失。

2. 双损失组合的合理性

   • 通过combined_loss = ce_loss + self.gamma * additional_loss实现的线性加权组合，是多任务损失融合的标准方案。gamma参数可以灵活控制额外损失对模型训练的影响程度。
   • 梯度基于最终的combined_loss计算，确保模型权重同时受到交叉熵、正则化、额外损失的共同约束，完全匹配你的需求。

代码优化建议 #

• 可以在train_step和test_step的返回字典中加入总损失的监控，方便观察训练过程：

  # train_step返回时添加
  return {
      **{m.name: m.result() for m in self.metrics},
      "combined_loss": combined_loss,
      "ce_loss": ce_loss,
      "additional_loss": additional_loss
  }
  

• test_step中可以保存验证时的损失结果，避免仅监控指标而忽略损失变化：

  def test_step(self, data):
      x, y = data
      y_pred = self.model(x, training=False)
      val_ce_loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
      val_additional_loss = self.another_loss(y, y_pred)
      val_combined_loss = val_ce_loss + self.gamma * val_additional_loss
      self.compiled_metrics.update_state(y, y_pred)
      return {
          **{m.name: m.result() for m in self.metrics},
          "val_combined_loss": val_combined_loss,
          "val_ce_loss": val_ce_loss,
          "val_additional_loss": val_additional_loss
      }
  

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