如何缓存并恢复tf.function编译的计算图？适配ExtensionType场景

一、直接缓存ConcreteFunction，绕过SavedModel的自动追踪 #

SavedModel对ExtensionType支持有限，且默认会触发重编译，直接保存tf.function生成的ConcreteFunction是更可靠的方式：

• 收集已追踪的ConcreteFunction：
  调用tf.function后，通过get_concrete_function(*args)获取对应输入签名的具体函数，或遍历func._concrete_functions获取所有已追踪版本（注：_concrete_functions是内部API，仅当前场景下推荐使用）。
  示例代码：

  import tensorflow as tf
  
  # 定义ExtensionType
  class MyType(tf.experimental.ExtensionType):
      value: tf.Tensor
  
  # 核心计算函数
  @tf.function(reduce_retracing=True)
  def core_func(inputs: MyType) -> MyType:
      return MyType(value=inputs.value * 2)
  
  # 触发追踪，生成ConcreteFunction
  test_input = MyType(value=tf.constant(1.0))
  concrete_func = core_func.get_concrete_function(test_input)
  
  # 保存到磁盘
  tf.saved_model.save(concrete_func, "./cached_func")
  

• 加载并直接使用：
  加载后直接调用ConcreteFunction，完全跳过重编译步骤：

  loaded_func = tf.saved_model.load("./cached_func")
  result = loaded_func(test_input)
  print(result.value.numpy())  # 输出2.0
  

  若存在多个retraced版本（不同输入形状/类型），需逐个保存对应的ConcreteFunction，或打包到自定义模块统一管理。

二、修复ExtensionType的SavedModel序列化错误 #

ExtensionType默认无法被SavedModel序列化，需手动添加序列化逻辑：

class MyType(tf.experimental.ExtensionType):
    value: tf.Tensor

    def _serialize(self):
        return {"value": self.value}

    @classmethod
    def _deserialize(cls, data):
        return cls(value=data["value"])

添加_serialize和_deserialize方法后，SavedModel可正确处理该类型的输入输出，避免ValueError。

三、跨会话无重编译的关键注意事项 #

• 加载时的输入签名必须与保存ConcreteFunction时完全一致（包括ExtensionType结构、张量形状/ dtype），否则仍会触发重编译。
• 直接调用加载后的ConcreteFunction，不要通过原tf.function装饰的函数间接调用，避免触发新的追踪。
• 批量管理多个ConcreteFunction可使用自定义tf.Module：

  class FuncCache(tf.Module):
      def __init__(self):
          self.func1 = concrete_func1
          self.func2 = concrete_func2
  
  # 保存模块
  tf.saved_model.save(FuncCache(), "./cached_funcs")
  
  # 加载调用
  loaded_cache = tf.saved_model.load("./cached_funcs")
  result = loaded_cache.func1(test_input)
  

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