如何改造TensorFlow实现的BCQ模型以支持SavedModel及TFLite转换？

错误原因 #

你遇到的ValueError是因为tf.saved_model.save()要求传入的对象必须是TensorFlow可追踪对象（继承自TrackableBase，比如tf.Module、tf.keras.Model），而你的BCQ类只是普通Python类，没有实现TensorFlow的状态追踪机制，所以无法被识别。

函数式编写模型与tf.Module类的差异 #

• 函数式编写：直接用TensorFlow运算API构建计算逻辑，变量通常分散在函数内部或手动创建，没有统一的状态管理。TensorFlow无法自动识别模型的结构和可训练参数，只能通过tf.compat.v1.train.Saver()保存权重，无法导出完整模型结构。
• tf.Module类编写：tf.Module是TensorFlow用于管理模型状态和结构的核心基类，所有可追踪对象（包括tf.keras.Model）都继承自它。它会自动收集类内部的tf.Variable和子模块，让TensorFlow能完整识别模型的结构、参数和计算逻辑，天然支持SavedModel导出、检查点管理和TFLite转换。

代码改造步骤 #

1. 让BCQ类继承tf.Module #

修改类定义，使其成为tf.Module的子类，同时将所有可训练参数转为类属性并使用tf.Variable定义：

import tensorflow as tf

class BCQ(tf.Module):
    def __init__(self, state_dim, action_dim, max_action, ...):
        super().__init__()
        self.state_dim = state_dim
        self.action_dim = action_dim
        # 示例：初始化Q网络的可训练参数
        self.q1_dense1_weights = tf.Variable(tf.random.normal([state_dim + action_dim, 256]), name="q1_dense1_weights")
        self.q1_dense1_bias = tf.Variable(tf.zeros([256]), name="q1_dense1_bias")
        # 其他网络参数（如actor网络、目标网络等）同理，全部转为类的属性
        ...

2. 封装前向推理逻辑为@tf.function装饰的方法 #

定义用于模型推理的方法（推荐用__call__），用@tf.function装饰并指定输入签名，确保TensorFlow能编译计算图并导出清晰的模型签名：

@tf.function(input_signature=[tf.TensorSpec(shape=[None, self.state_dim], dtype=tf.float32)])
def __call__(self, state):
    # 在这里实现BCQ模型的完整前向计算逻辑，使用类中定义的self.xxx参数
    # 示例：Q网络前向计算
    h1 = tf.matmul(tf.concat([state, action], axis=1), self.q1_dense1_weights) + self.q1_dense1_bias
    h1 = tf.nn.relu(h1)
    ...
    return predicted_q_values

3. 保存为SavedModel格式 #

改造完成后，直接使用tf.saved_model.save()保存完整模型：

bcq_model = BCQ(state_dim, action_dim, max_action, ...)
# 先完成模型训练...
tf.saved_model.save(bcq_model, "./bcq_saved_model")

4. 转换为TFLite格式 #

通过SavedModel路径直接完成转换：

converter = tf.lite.TFLiteConverter.from_saved_model("./bcq_saved_model")
tflite_model = converter.convert()
# 保存TFLite模型文件
with open("./bcq_model.tflite", "wb") as f:
    f.write(tflite_model)

注意事项 #

• 确保所有可训练参数都用tf.Variable定义并作为类属性，避免在函数内部临时创建变量，否则TensorFlow无法追踪状态。
• 若原有代码使用tf.compat.v1API，建议逐步迁移到TensorFlow 2.x原生API，减少兼容性问题。
• 训练逻辑无需大幅修改，只需保证训练时更新的是类的tf.Variable属性即可。

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