tfa.optimizers.MultiOptimizer - TypeError: 'Not JSON Serializable:'

Question

我正在尝试使用 tfa.optimizers.MultiOptimizer()。我根据文档 (https://www.tensorflow.org/addons/api_docs/python/tfa/optimizers/MultiOptimizer) 做了所有事情，但出现以下错误：

TypeError: ('Not JSON Serializable:', )

以下是重现错误的最小工作示例，只需复制并粘贴即可。当第一个 epoch 完成并且回调尝试保存模型时会发生错误。

##############################################################################

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

##############################################################################

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()]
    
    optimizers_and_layers = [(optimizers[0], model.layers[:5]), (optimizers[1], model.layers[5:])]
    
    optimizer = tfa.optimizers.MultiOptimizer(optimizers_and_layers)
    
    model.compile(optimizer=optimizer, loss='mse', metrics='mse')

    test = tf.keras.optimizers.serialize(optimizer)

    return model

##############################################################################

input_data =  np.arange( 0, 10000, 1).reshape(10000,1)
target_data = np.arange(-10000, 0, 1).reshape(10000,1)

model = build_model_1()

model_checkpoint = ke.callbacks.ModelCheckpoint('best_model.h5',
                                                monitor='val_mse',
                                                mode='min',
                                                save_best_only=True,
                                                verbose=1)

training_history = model.fit(x = input_data,
                             y = target_data,
                             validation_split = 0.2,
                             epochs = 5,
                             verbose = 1,
                             callbacks = [model_checkpoint])
    
##############################################################################

Answer 1

当保存 完整 Keras 模型（在 .h5 文件中具有自己的结构）时，tf.keras.Model 对象被完全序列化为 JSON：这意味着模型的每个属性都应该是 JSON 可序列化的。

注意：tf.Tensor 不是 JSON 可序列化的。

当使用来自tfa 的这个多重优化器时，您将向模型添加属性，JSON 序列化器将尝试（并且失败）序列化。

特别是 gv 这个属性我认为它来自使用的自定义优化器。

'gv': [(<tf.Tensor 'gradient_tape/model/dense/Tensordot/MatMul/MatMul:0' shape=(1, 1) dtype=float32>, <tf.Variable 'dense/kernel:0' shape=(1, 1) dtype=float32, numpy=array([[-0.55191684]], dtype=float32)>), (<tf.Tensor 'gradient_tape/model/dense/BiasAdd/BiasAddGrad:0' shape=(1,) dtype=float32>, <tf.Variable 'dense/bias:0' shape=(1,) dtype=float32, numpy=array([-0.23444518], dtype=float32)>)]},

所有这些tf.Tensor 都不是 JSON 可序列化的，这就是它失败的原因。

唯一的选择是不保存模型完全（及其所有属性，应定义为 Keras 层，但在这种情况下是不可能的），而只保存模型参数。

简而言之，如果您将 save_weights_only=True 添加到回调中，您的训练（和权重检查点）将正常工作。

model_checkpoint = ke.callbacks.ModelCheckpoint(
    "best_model.h5",
    monitor="val_mse",
    mode="min",
    save_best_only=True,
    verbose=1,
    save_weights_only=True,
)