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

Issue

I’m trying to use tfa.optimizers.MultiOptimizer(). I did everything according to the docs (https://www.tensorflow.org/addons/api_docs/python/tfa/optimizers/MultiOptimizer) yet I’m getting the following error:

TypeError: (‘Not JSON Serializable:’, <tf.Tensor ‘gradient_tape/model_80/dense_3/Tensordot/MatMul/MatMul:0’ shape=(1, 1) dtype=float32>)

Below is a minimal, working example that reproduces the error, just copy and paste it. The error occurs when the first epoch is finished and the callback trys to save the model.

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

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])
    
##############################################################################

Solution

When saving a complete Keras model (with its own structure in the .h5 file) the tf.keras.Model object is completely serialized as a JSON: this means that every property of the model should be JSON serializable.

NOTE: tf.Tensor are NOT JSON serializable.

When using this multi optimizer from tfa you’re adding properties to the model that the JSON serializer will try (and fail) to serialize.

In particular there’s this attribute gv that I think it comes from the custom optimizer used.

'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)>)]},

All this tf.Tensor are not JSON serializable, that’s why it fails.

The only option is to do not save the model completely (with all its attributes, which should be defined as Keras layers, but in this case is not possible) but saving only the model parameters.

In short, if you add the save_weights_only=True to the callback your training (and checkpoint of the weights) will work fine.

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

Answered By – nessuno

This Answer collected from stackoverflow, is licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0