Saving and loading Tensorflow probability layers inside Keras sequential module

Question

I am using Tensorflow probability layers inside Keras sequentials. However, saving the model as json and then loading it throws and exception. I am using custom_objects to be able to load custom layers. Here is the minimalistic code to reproduce the error.

import tensorflow_probability as tfp

tfk = tf.keras
tfkl = tf.keras.layers
tfpl = tfp.layers

original_dim = 20
latent_dim = 2
model = tfk.Sequential([
    tfkl.InputLayer(input_shape=original_dim),
    tfkl.Dense(10, activation=tf.nn.leaky_relu),
    tfkl.Dense(tfpl.MultivariateNormalTriL.params_size(latent_dim), activation=None),
    tfpl.MultivariateNormalTriL(latent_dim)
])

model_json = model.to_json()
with open("model.json", "w") as json_file:
    json_file.write(model_json)



loaded_model = tfk.models.model_from_json(
    open('model.json').read(),
    custom_objects={
        'leaky_relu': tf.nn.leaky_relu, 
        'MultivariateNormalTriL': tfpl.MultivariateNormalTriL
    }
)

I get the following exception:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-26-bbbeffd9e4be> in <module>
      3     custom_objects={
      4         'leaky_relu': tf.nn.leaky_relu,
----> 5         'MultivariateNormalTriL': tfpl.MultivariateNormalTriL
      6     }
      7 )

//anaconda3/envs/dl-env/lib/python3.7/site-packages/tensorflow/python/keras/saving/model_config.py in model_from_json(json_string, custom_objects)
     94   config = json.loads(json_string)
     95   from tensorflow.python.keras.layers import deserialize  # pylint: disable=g-import-not-at-top
---> 96   return deserialize(config, custom_objects=custom_objects)

//anaconda3/envs/dl-env/lib/python3.7/site-packages/tensorflow/python/keras/layers/serialization.py in deserialize(config, custom_objects)
     87       module_objects=globs,
     88       custom_objects=custom_objects,
---> 89       printable_module_name='layer')

//anaconda3/envs/dl-env/lib/python3.7/site-packages/tensorflow/python/keras/utils/generic_utils.py in deserialize_keras_object(identifier, module_objects, custom_objects, printable_module_name)
    190             custom_objects=dict(
    191                 list(_GLOBAL_CUSTOM_OBJECTS.items()) +
--> 192                 list(custom_objects.items())))
    193       with CustomObjectScope(custom_objects):
    194         return cls.from_config(cls_config)

//anaconda3/envs/dl-env/lib/python3.7/site-packages/tensorflow/python/keras/engine/sequential.py in from_config(cls, config, custom_objects)
    350     for layer_config in layer_configs:
    351       layer = layer_module.deserialize(layer_config,
--> 352                                        custom_objects=custom_objects)
    353       model.add(layer)
    354     if not model.inputs and build_input_shape:

//anaconda3/envs/dl-env/lib/python3.7/site-packages/tensorflow/python/keras/layers/serialization.py in deserialize(config, custom_objects)
     87       module_objects=globs,
     88       custom_objects=custom_objects,
---> 89       printable_module_name='layer')

//anaconda3/envs/dl-env/lib/python3.7/site-packages/tensorflow/python/keras/utils/generic_utils.py in deserialize_keras_object(identifier, module_objects, custom_objects, printable_module_name)
    190             custom_objects=dict(
    191                 list(_GLOBAL_CUSTOM_OBJECTS.items()) +
--> 192                 list(custom_objects.items())))
    193       with CustomObjectScope(custom_objects):
    194         return cls.from_config(cls_config)

//anaconda3/envs/dl-env/lib/python3.7/site-packages/tensorflow_probability/python/layers/distribution_layer.py in from_config(cls, config, custom_objects)
    875             config['arguments'][key] = np.array(arg_dict['value'])
    876 
--> 877     return cls(**config)
    878 
    879   @classmethod

TypeError: __init__() missing 1 required positional argument: 'event_size'

score 1 · Answer 1 · answered Jul 02 '20 at 06:29

I had the same problem. I solved it by adding this to custom_objects

    def MultivariateNormalTriL_loader(latent_dim):
        def load_MultivariateNormalTriL(name, trainable, type, 
                                        function, function_type, module, 
                                        output_shape, output_shape_type,  
                                        output_shape_module, arguments, 
                                        make_distribution_fn, convert_to_tensor_fn):
            return tfp.layers.MultivariateNormalTriL(latent_dim, name=name, 
                                          trainable=trainable, dtype=dtype, 
                                          convert_to_tensor_fn=convert_to_tensor_fn)
        return load_MultivariateNormalTriL
    # Use the latent_dim here
    custom_objects['MultivariateNormalTriL'] = MultivariateNormalTriL_loader(latent_dim)

I'm not sure about which arguments are necessary, but these worked for me.

Thanks for your answer. But this means we need to know part of the model parameter prior to loading. I was wondering if we can save `latent_dim` as part of model parameters as well. — Hamed, Jul 03 '20 at 07:24

score -1 · Answer 2 · answered Oct 09 '19 at 06:26

-1

See if the following load method works:

loaded_model = tfk.models.model_from_json(
    open('model.json').read(),
    custom_objects={
        'leaky_relu': tf.nn.leaky_relu, 
        'MultivariateNormalTriL': tfpl.MultivariateNormalTriL.params_size(latent_dim)
    }
)

answered Oct 09 '19 at 06:26

Mohammad

1,006
2
15
29

Not surprisingly this won't work. `params_size` returns an integer. We have to provide class types inside `custom_objects`. – Hamed Oct 09 '19 at 08:43

Saving and loading Tensorflow probability layers inside Keras sequential module

2 Answers2