tf.function装饰子类模型call方法后模型无法收敛的问题

我在复杂模型中使用全连接神经网络作为函数逼近器，使用Functional API构建模型时可正常收敛并得到合理结果，代码如下：

inputs = keras.Input(shape=(setupDict['inputShape'],), name='input')
x1 = keras.layers.Dense(setupDict['layerNodes'][0], activation='relu', kernel_initializer='he_normal', name='hidden1')(inputs)
x2 = keras.layers.Dense(setupDict['layerNodes'][1], activation='relu', kernel_initializer='he_normal', name='hidden2')(x1)
output1 = keras.layers.Dense(1, kernel_initializer='he_normal', name='outputD')(x2)
output1 = LogisticActivation(0.15, 1.5, 5, 'debt')(output1)
output2 = keras.layers.Dense(1, kernel_initializer='he_normal', name='outputS')(x2)
output2 = LogisticActivation(0.005, 6, 10, 'stab')(output2)
output3 = keras.layers.Dense(1, activation='softplus', kernel_initializer='he_normal', bias_initializer='ones', name='outputV')(x2)
outputs = keras.layers.concatenate([output1, output2, output3], name='concatAll')
# Define the model
model = keras.Model(inputs=inputs,
                    outputs=outputs)

因存在更复杂的使用场景，我将其改写为模型子类化实现，并尝试通过tf.function启用图执行。改写后的模型代码如下：

class defaultFreeModel(keras.Model):
    def __init__(self, econDict, setupDict, GHDict, **kwargs):
        super().__init__(**kwargs)
        self.econDict = econDict
        self.setupDict = setupDict
        self.GHDict = GHDict
        # Layers
        self.hidden1 = keras.layers.Dense(setupDict['layerNodes'][0], activation='relu', kernel_initializer='he_normal', name='hidden1')
        self.hidden2 = keras.layers.Dense(setupDict['layerNodes'][1], activation='relu', kernel_initializer='he_normal', name='hidden2')
        self.output1tmp = keras.layers.Dense(1, kernel_initializer='he_normal', name='outputD')
        self.output1 = LogisticActivation(0.15, 1.5, 5, 'debt')
        self.output2tmp = keras.layers.Dense(1, kernel_initializer='he_normal', name='outputS')
        self.output2 = LogisticActivation(0.005, 6, 10, 'stab')
        self.output3 = keras.layers.Dense(1, activation='softplus', kernel_initializer='he_normal', bias_initializer='ones',
                                     name='outputV')
        self.outputs = keras.layers.Concatenate()


    def call(self, inputs):
        x = self.hidden1(inputs)
        x = self.hidden2(x)
        x1 = self.output1tmp(x)
        x1 = self.output1(x1)
        x2 = self.output2tmp(x)
        x2 = self.output2(x2)
        x3 = self.output3(x)
        return self.outputs([x1, x2, x3])

    def compile(self, optimizer, loss_fn):
        super().compile()
        self.optimizer = optimizer
        self.loss_fn = loss_fn

未给call方法添加tf.function装饰器时，训练循环运行正常，模型可正常收敛。但给call方法添加该装饰器以加速执行后，模型不再收敛，损失值在某一数值附近波动（因初始化器不同，每次运行的数值不同）。我还尝试给梯度函数也添加tf.function装饰器，虽大幅提升了计算速度，但模型仍无法收敛。自定义LogisticActivation层的代码如下：

import tensorflow as tf
import keras
class LogisticActivation(keras.layers.Layer):
    def __init__(self, level, slope, upper, name):
        super().__init__()
        self.slope = tf.constant(slope, dtype=tf.float32, name=name + "_slope")
        self.level = tf.constant(level, dtype=tf.float32, name=name + "_level")
        self.upper = tf.constant(upper, dtype=tf.float32, name=name + "_upper")

    @tf.function
    def call(self, inputs):
        return tf.divide(tf.multiply(self.level, self.upper), tf.constant(1, dtype=tf.float32)+tf.exp(-self.slope*(inputs-self.level)))

    def get_config(self):
        return {'level': self.level, 'slope': self.slope, 'upper': self.upper}

请问导致这一问题的原因可能是什么？

• 自定义compile方法的实现错误
  你重写的compile方法没有正确调用父类的核心训练逻辑：

  def compile(self, optimizer, loss_fn):
      super().compile()  # 未传入optimizer和loss参数，导致Keras内部训练机制未初始化
      self.optimizer = optimizer
      self.loss_fn = loss_fn
  

  Keras的Model.compile()需要接收optimizer、loss等参数来初始化梯度跟踪、参数更新等核心流程。仅手动保存optimizer和loss_fn，会导致图模式下模型无法正确处理参数更新，最终表现为损失不收敛。正确写法应为：

  def compile(self, optimizer, loss_fn, **kwargs):
      super().compile(optimizer=optimizer, loss=loss_fn, **kwargs)
  

  或者直接删除自定义的compile方法，使用父类默认实现即可。

• 嵌套tf.function的冲突
  自定义LogisticActivation层的call方法已添加@tf.function装饰器，同时又给模型的call方法添加该装饰器，形成嵌套图编译。这种嵌套可能导致梯度流被意外截断，或者图优化逻辑干扰梯度计算的正确性，使得模型参数无法得到有效更新。建议移除自定义层call方法上的@tf.function，模型的call被装饰后，层的前向逻辑会自动纳入图执行，无需单独装饰。

• Concatenate层的调用方式隐患
  在子类模型中，你将Concatenate层实例化为self.outputs并在call中调用。虽然写法本身合法，但在图模式下，若输入张量的batch size动态变化，可能引发形状推断错误，进而影响前向传播和梯度计算。可以改用更直接的tf.concat函数替代：

  return tf.concat([x1, x2, x3], axis=-1)
  

• 自定义层常量定义的不规范
  在LogisticActivation层的__init__中，你使用tf.constant保存固定参数，这会导致这些常量被直接嵌入计算图，同时get_config返回的是张量对象而非可序列化的Python数值，可能引发图编译或模型序列化的潜在问题。更规范的写法是先保存Python数值，在call中再转换为张量：

  class LogisticActivation(keras.layers.Layer):
      def __init__(self, level, slope, upper, name):
          super().__init__()
          self.slope = slope
          self.level = level
          self.upper = upper
          self.name = name
  
      def call(self, inputs):
          slope = tf.constant(self.slope, dtype=tf.float32, name=f"{self.name}_slope")
          level = tf.constant(self.level, dtype=tf.float32, name=f"{self.name}_level")
          upper = tf.constant(self.upper, dtype=tf.float32, name=f"{self.name}_upper")
          return tf.divide(tf.multiply(level, upper), 1.0 + tf.exp(-slope*(inputs-level)))
  
      def get_config(self):
          return {'level': self.level, 'slope': self.slope, 'upper': self.upper, 'name': self.name}
  

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