Keras自定义train_step函数对人脸超分辨率模型训练的影响

针对你复现DIDnet人脸超分辨率模型时遇到的问题——仅修改train_step的返回参数（从仅返回损失改为返回所有指标）就导致训练结果大幅下降，以下是可能的技术原因及对应解决方案：

可能的原因 #

• 梯度计算被指标干扰
  自定义train_step中，如果在计算指标时没有隔离梯度，指标计算过程中用到的模型输出张量会被GradientTape追踪，导致反向传播时不仅基于损失计算梯度，还会混入指标计算的梯度，直接改变参数更新的方向，偏离原本的损失优化目标。

• 张量生命周期与自动微分异常
  返回多个指标时，模型输出张量可能被多次引用，触发Keras自动微分机制的隐性异常。比如部分张量的追踪状态被修改，导致梯度计算的路径出现错误，最终影响参数更新的正确性。

• Keras训练循环的隐性处理逻辑
  Keras对train_step返回的所有张量会进行内部追踪，如果返回的指标包含与可训练变量的依赖关系，可能会被误判为需要优化的额外目标，或者干扰损失梯度的计算流程，导致参数更新混乱。

• 指标计算引发的数值/形状异常
  计算指标时可能不经意间修改了张量的数据类型（如float32转float64）或形状，导致后续损失计算、梯度更新出现精度损失或维度不匹配，最终影响训练效果。

解决方案 #

• 隔离指标计算的梯度
  在计算评价指标时，用tf.stop_gradient()包裹模型输出或相关张量，确保指标计算不参与反向传播。示例代码：

  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)
      
      # 更新模型参数
      gradients = tape.gradient(loss, self.trainable_variables)
      self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
      
      # 隔离梯度计算指标
      with tf.stop_gradient():
          self.compiled_metrics.update_state(y, y_pred)
      
      # 返回所有指标与损失
      return {m.name: m.result() for m in self.metrics}
  

• 分离损失与指标计算逻辑
  确保损失计算完全在GradientTape上下文内，指标计算在上下文外，且指标计算仅使用已脱离梯度追踪的张量，避免干扰损失的梯度计算路径。

• 检查自定义指标实现
  确认自定义指标的代码中没有修改模型的可训练变量，也没有产生额外的梯度依赖（比如避免在指标中使用可训练层的输出而不隔离梯度）。

• 验证返回张量的一致性
  打印返回的损失张量（两种返回方式下）的形状、数据类型和数值，确保仅返回损失时和返回所有指标时的损失值初始状态一致，排除隐性的数值异常。

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