升级TensorFlow 2.16后Keras自定义模型@tf.function报位置参数过多错误

问题根源 #

报错的核心原因是自定义Model类中错误使用了InputLayer。InputLayer是为Keras函数式API设计的，在子类化tf.keras.Model时，完全不需要显式创建并调用它。旧版TensorFlow的tf.function对函数签名检查不严格，这个错误被掩盖；但TensorFlow 2.12及以后版本改用Python inspect直接解析函数签名，原本的非法调用触发了too many positional arguments错误。

修复步骤 #

1. 移除__init__方法中self.input_layer的定义
2. 修改call方法，直接将输入张量传给第一个隐藏层，删除调用self.input_layer的代码

修复后的完整Model代码 #

class Model(tf.keras.Model):
    """
    Subclassing a multi-layered NN using Keras from Tensorflow
    """

    def __init__(self, num_states, hidden_units, num_actions):
        super(Model, self).__init__()
        self.hidden_layers = []

        for hidden_unit in hidden_units:
            self.hidden_layers.append(
                tf.keras.layers.Dense(hidden_unit, activation="tanh")
            )

        self.output_layer = tf.keras.layers.Dense(num_actions, activation="linear")

    @tf.function
    def call(self, inputs):
        x = inputs
        for layer in self.hidden_layers:
            x = layer(x)
        output = self.output_layer(x)
        return output

额外优化（非报错必需） #

select_action方法中重复调用np.atleast_2d可以简化，避免生成多余的维度：

def select_action(self, state, policy_net):
    rate = self.strategy.get_exploration_rate(self.current_step)
    self.current_step += 1

    if rate > random.random():
        return random.randrange(self.num_actions)
    else:
        state = np.atleast_2d(state).astype("float32")
        return np.argmax(policy_net(state))

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