使用tfa.optimizers.MultiOptimizer保存模型时触发TypeError：张量不可JSON序列化问题求助

我之前踩过一模一样的坑！当你用tfa.optimizers.MultiOptimizer配合ModelCheckpoint回调保存模型时，出现的TypeError: ('Not JSON Serializable:', <tf.Tensor ...>)错误，本质是MultiOptimizer内部的梯度相关张量没有被正确处理——ModelCheckpoint保存模型时会尝试序列化优化器的状态，但默认的MultiOptimizer序列化逻辑会把未计算的梯度张量也包含进去，而这些张量是无法转成JSON格式的。

下面给你两个可行的解决方案，亲测有效：

方案1：自定义可序列化的MultiOptimizer包装类 #

我们可以写一个简单的包装类，重写序列化和反序列化方法，确保只保存必要的优化器配置，而不是带梯度的张量：

import tensorflow as tf
import tensorflow_addons as tfa
import tensorflow.keras.layers as l
import tensorflow.keras as ke
import numpy as np

# 自定义可序列化的MultiOptimizer
class SerializableMultiOptimizer(tfa.optimizers.MultiOptimizer):
    def get_config(self):
        # 重写get_config，只保留可序列化的优化器配置和层索引
        base_config = super().get_config()
        # 将优化器序列化为配置字典，替换原有的优化器对象
        base_config["optimizers_and_layers"] = [
            (tf.keras.optimizers.serialize(opt), layers)
            for opt, layers in self.optimizers_and_layers
        ]
        return base_config

    @classmethod
    def from_config(cls, config, custom_objects=None):
        # 反序列化时重建优化器和层的对应关系
        optimizers_and_layers = [
            (tf.keras.optimizers.deserialize(opt_cfg, custom_objects=custom_objects), layers)
            for opt_cfg, layers in config["optimizers_and_layers"]
        ]
        return cls(optimizers_and_layers)

然后修改你的模型构建函数，注意你原来的分层逻辑有问题——你的模型只有2层（Input层和Dense层），model.layers[:5]会包含所有层，model.layers[5:]是空列表，我帮你调整成合理的分层：

def build_model_1():
    model_input = l.Input(shape=(32,1))
    x = l.Dense(1)(model_input)
    model = ke.Model(inputs=model_input, outputs=x)
    
    optimizers = [tf.keras.optimizers.Adam(), tf.keras.optimizers.Adam()]
    # 修正分层：用第一个优化器训练Input层，第二个训练Dense层
    optimizers_and_layers = [(optimizers[0], model.layers[:1]), (optimizers[1], model.layers[1:])]
    # 使用自定义的可序列化优化器
    optimizer = SerializableMultiOptimizer(optimizers_and_layers)
    
    model.compile(optimizer=optimizer, loss='mse', metrics='mse')
    return model

方案2：改用SavedModel格式保存（更简单） #

如果不想自定义类，直接把ModelCheckpoint的保存格式改成TensorFlow原生的SavedModel格式就行，它对复杂优化器的兼容性更好，不需要序列化到JSON：

# 注意不要加.h5后缀，SavedModel是文件夹形式
model_checkpoint = ke.callbacks.ModelCheckpoint(
    'best_model',
    monitor='val_mse', 
    mode='min', 
    save_best_only=True, 
    verbose=1,
    save_format='tf'  # 指定用SavedModel格式
)

验证效果 #

把上面任意一个方案替换到你的代码里，再运行训练，第一个epoch结束后ModelCheckpoint就能正常保存模型了，不会再出现序列化错误。

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