TensorFlow 2.4用tf.feature_columns建模 加载时触发tf.function重追踪警告

警告是否为流程预期表现 #

该警告是TensorFlow 2.4.0版本下，tf.feature_columns组件与SavedModel保存、加载流程联动产生的已知常见现象，不属于你的业务代码编写错误。
具体触发逻辑：tf.feature_columns是TensorFlow 1.x遗留的兼容类组件，在2.4.0版本中，其内置的DenseFeatures层在序列化为SavedModel时，会自动生成无显式装饰的隐式tf.function处理逻辑。加载模型时，该内置函数默认采用严格的形状匹配规则，且2.4.0版本未对这类内置函数做默认的形状松弛优化，即使你没有自定义@tf.function，也会在模型加载的自动预热调用过程中触发多次重追踪，最终抛出该警告。
由于触发重追踪的是TensorFlow内部自动生成的隐式函数，你没有权限直接给这类函数手动添加experimental_relax_shapes=True参数来优化，因此无法通过常规的自定义函数配置解决该问题。

彻底消除警告的可行方案 #

• 替换模型加载接口
  将测试代码中的tf.saved_model.load(temp_dir)替换为tf.keras.models.load_model(temp_dir)。二者的核心差异是：tf.saved_model.load会重新解析SavedModel的计算图、重新生成隐式函数的追踪逻辑；而tf.keras.models.load_model会直接复用模型保存时已经完成的函数追踪结果，从根源减少不必要的重追踪。
• 调整输入层形状定义，固定动态形状特征
  将原Input层定义：

  inputs[col_name] = tf.keras.layers.Input(name=col_name, shape=(1,), dtype=tf.dtypes.string)
  

  修改为：

  inputs[col_name] = tf.keras.layers.Input(name=col_name, batch_shape=(None, 1), dtype=tf.dtypes.string)
  

  该修改会在函数追踪阶段固定输入的非batch维度形状，避免因形状推导不一致触发的重追踪。
• 替换遗留的tf.feature_columns为原生Keras预处理层
  tf.feature_columns目前已进入维护模式，官方更推荐使用2.x版本原生的Keras预处理层实现相同的特征处理逻辑，可从根源规避该类警告。原有特征列构建逻辑可直接替换为：

  # 原有遗留逻辑（触发警告的根源）
  # categorical_column = tf.feature_column.categorical_column_with_vocabulary_list(col_name, ['A', 'B', 'C', 'D'])
  # feature_columns[col_name] = tf.feature_column.indicator_column(categorical_column=categorical_column)
  
  # 替换为原生Keras层逻辑
  vocab_layer = tf.keras.layers.StringLookup(
      vocabulary=['A', 'B', 'C', 'D'],
      output_mode='one_hot',
      name=f'{col_name}_vocab'
  )
  

  后续直接将该层挂接到Keras计算图中即可。这类原生层的SavedModel序列化、加载逻辑经过专门优化，不会触发不必要的重追踪警告。
• 临时兜底方案：过滤警告输出
  如果以上调整受业务场景限制无法实施，可通过调整TensorFlow日志级别，在代码开头添加如下配置屏蔽该类警告输出：

  import os
  os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
  import tensorflow as tf
  tf.get_logger().setLevel('ERROR')
  

  注意：该方案仅屏蔽警告输出，并没有从根源解决重追踪带来的性能损耗，仅适合临时使用。

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