如何解决TensorFlow获取concrete函数时出现的IndexError报错

这个索引越界错误是tf.function包装Keras自定义层的方式错误，导致参数解析错位引发的：

• Keras层实例的__call__方法不是纯计算函数，它内部会自动处理层构建（build）、训练/推理模式切换、掩码传递等额外逻辑，除了用户传入的inputs参数外，还隐式包含training、mask等可选入参。
• 原代码直接将层实例inc传入tf.function做包装，又没有提前完成层的build初始化，后续调用get_concrete_function只传入了输入的TensorSpec时，TensorFlow在追踪函数的过程中会把层内部维护的参数列表、隐式入参和显式传入的输入签名混排，最终按索引取参数时超出列表长度，触发IndexError: list index out of range。

按以下步骤排查修复即可：

• 第一步先验证层本身逻辑正确性
  单独运行层的前向传播，排除自定义层本身的代码错误：

  inc = Inc()
  # 测试层前向传播
  print(inc(tf.constant([[1.0, 2.0, 3.0]])))
  

  如果正常输出tf.Tensor([[2. 3. 4.]], shape=(1, 3), dtype=float32)，说明层逻辑无问题，故障点在tf.function包装和concrete函数提取环节。
• 修复方式一：包装层的call方法，提前完成层构建
  不要直接包装层实例，改为包装层的call方法，同时提前调用build传入输入形状完成层初始化，避免参数追踪错位：

  import tensorflow as tf
  
  class Inc(tf.keras.layers.Layer):
      def call(self, inputs):
          return inputs + 1
  inc = Inc()
  # 提前构建层，完成参数初始化
  inc.build(input_shape=(1, 3))
  # 包装纯计算逻辑的call方法，而非层实例本身
  inc_g2 = tf.function(inc.call)
  
  concrete_fun = inc_g2.get_concrete_function(tf.TensorSpec(shape=(1, 3), dtype=tf.float32))
  print(concrete_fun(tf.constant([[1.0, 2.0, 3.0]])))
  

• 修复方式二：包装层实例时显式指定输入签名
  如果需要直接包装层实例，要在tf.function初始化时通过input_signature参数明确传入参数格式，避免自动追踪时的参数解析错误：

  import tensorflow as tf
  
  class Inc(tf.keras.layers.Layer):
      def call(self, inputs):
          return inputs + 1
  inc = Inc()
  # 包装实例时直接明确输入签名
  inc_g2 = tf.function(inc, input_signature=[tf.TensorSpec(shape=(1, 3), dtype=tf.float32)])
  # 已经指定过输入签名，不需要再重复传TensorSpec
  concrete_fun = inc_g2.get_concrete_function()
  print(concrete_fun(tf.constant([[1.0, 2.0, 3.0]])))
  

注意：提取concrete函数时传入的TensorSpec形状、数据类型，必须和实际调用时传入的张量严格匹配，否则也会触发追踪错误。

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