如何保存含用作配置参数的tf.Variables的TensorFlow模型？

问题根源 #

你遇到的序列化错误，是因为直接将tf.Variable作为Keras模型的属性存储。模型保存时，Keras会尝试将所有模型属性序列化为JSON格式，但tf.Variable不属于可JSON序列化的类型，因此触发报错。

结合你的场景（无训练需求、需实时调参、要部署到TFLite），最合理的解决方案是将可调参数作为模型的输入层，而非类内的tf.Variable属性。这种方式既符合Keras设计规范，又完美兼容TFLite部署，同时满足实时调参需求。

修正代码示例 #

1. 修改Sine层，支持多输入参数

import tensorflow as tf
import numpy as np

class Sine(tf.keras.layers.Layer):
    def __init__(self, *args, **kwargs):
        super(Sine, self).__init__(*args, **kwargs)
        self._twopi = tf.constant(np.pi * 2.0)
        
    def call(self, inputs):
        # 拆分输入：时间 + 四个可调参数
        time, scale, frequency, base, phase = inputs
        time = tf.cast(time, tf.float32)
        return scale * tf.sin(self._twopi * frequency * time + phase) + base

2. 重构模型，将参数作为输入层

class StupidModel:
    def __init__(self):
        self._model = self._build_model()

    def _build_model(self):
        # 定义多输入：时间序列 + 四个可调参数
        time_input = tf.keras.layers.Input(shape=(1,), name="time")
        amplitude_input = tf.keras.layers.Input(shape=(), name="amplitude")
        frequency_input = tf.keras.layers.Input(shape=(), name="frequency")
        base_input = tf.keras.layers.Input(shape=(), name="base")
        phase_input = tf.keras.layers.Input(shape=(), name="phase")
        
        # 传入Sine层计算输出
        out = Sine()([time_input, amplitude_input, frequency_input, base_input, phase_input])
        
        # 构建多输入模型
        model = tf.keras.Model(
            inputs=[time_input, amplitude_input, frequency_input, base_input, phase_input],
            outputs=out
        )
        return model

    def __call__(self, time, amplitude=1.0, frequency=0.5, base=0.0, phase=0.0):
        # 调用时可直接传入参数，默认值方便快速测试
        return self._model.predict([time, amplitude, frequency, base, phase])

3. 测试保存与调用

# 初始化模型
sm = StupidModel()
# 先调用一次构建模型结构
sm(tf.linspace(0.0, 1.0, 100))
# 保存模型（无序列化错误）
sm._model.save("foo.tf")

# 实时调参示例
time_steps = tf.linspace(0.0, 2.0, 200)
# 调整频率和振幅
result = sm(time_steps, frequency=0.8, amplitude=1.5)

方案优势 #

1. 解决序列化问题：所有可调参数作为输入层，Keras仅序列化输入层配置，无自定义属性导致的序列化错误。
2. 适配TFLite部署：多输入结构完全支持TFLite转换，移动端/网页端可直接传入参数值实现实时调参。
3. 灵活调参：不管是Python CLI还是部署后的应用，都能通过传入不同参数值快速调整刺激生成逻辑，无需重新生成或上传模型。

备选方案（不推荐，仅作参考） #

如果坚持使用tf.Variable，可通过重写模型的get_config方法实现自定义序列化，但该方案不适用于TFLite动态调参（Variable会被转为固定权重）：

class CustomModel(tf.keras.Model):
    def __init__(self, frequency, amplitude, base, phase, **kwargs):
        super().__init__(**kwargs)
        self._frequency = tf.Variable(frequency, name="frequency", dtype=tf.float32)
        self._amplitude = tf.Variable(amplitude, name="amplitude", dtype=tf.float32)
        self._base = tf.Variable(base, name="base", dtype=tf.float32)
        self._phase = tf.Variable(phase, name="phase", dtype=tf.float32)
        self.sine_layer = Sine()

    def call(self, inputs):
        time = inputs
        return self.sine_layer([time, self._amplitude, self._frequency, self._base, self._phase])

    def get_config(self):
        config = super().get_config()
        # 将Variable值转为普通数值保存
        config.update({
            "frequency": self._frequency.numpy(),
            "amplitude": self._amplitude.numpy(),
            "base": self._base.numpy(),
            "phase": self._phase.numpy()
        })
        return config

    @classmethod
    def from_config(cls, config):
        # 加载时从数值恢复Variable
        return cls(
            frequency=config["frequency"],
            amplitude=config["amplitude"],
            base=config["base"],
            phase=config["phase"]
        )

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