How to save a model with DenseVariational layer?

Question

I'm trying to build a model with DenseVariational layer so that it can report epistemic uncertainties. Something like https://www.tensorflow.org/probability/examples/Probabilistic_Layers_Regression#figure_3_epistemic_uncertainty

The model training works just fine and now I would like to save the model and load it in a production environment. However, when I tried model.save('path/model.h5'), I got

Layer DenseVariational has arguments in `__init__` and therefore must override `get_config`.

Then I added

class CustomVariational(tfp.layers.DenseVariational):
  def get_config(self):
        config = super().get_config().copy()
        config.update({
            'units': self.units,
            'make_posterior_fn': self._make_posterior_fn,
            'make_prior_fn': self._make_prior_fn
        })
        return config

but it failed with a new error

Unable to create link (name already exists)

Is DenseVariational layer for research only?

Answer 1

It's been almost 2 years, and the problem is still going on.

A workaround is to store only the weights:

tf.keras.Model.save_weights(filepath, overwrite=True)

Then, you can use the same model structure and load the weights.

For example:

# model
model = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=(input_shape,), name="input"),
    tfp.layers.DenseVariational(32, posterior_mean_field, prior_trainable),  # trainable
    tf.keras.layers.Dense(32, activation="relu"),
    tf.keras.layers.Dense(1)
])

# save weights after compiling and training your model
model.save_weights('model_weights.h5')

Initialize a new model with the same structure:

# different model, same weights
model2 = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=(input_shape,), name="input"),
    tfp.layers.DenseVariational(32, posterior_mean_field, prior_trainable),
    tf.keras.layers.Dense(32, activation="relu"),
    tf.keras.layers.Dense(1)
])

# load weights
model2.load_weights('model_weights.h5')

I hope this helps!

Answer 2

I think we can circumvent this problem by using the save_weights method.

Answer 3

When you add with tf.name_scope(...) to prior & posterior functions, it should be resolved, otherwise they end up with the same name for both tensors.

We're also fixing the example tutorial colab, should be online soon, thanks.

Update:

Instead of fixing it at the applications, we fixed it in the library instead: https://github.com/tensorflow/probability/commit/0ca065fb526b50ce38b68f7d5b803f02c78c8f16. Once it is updated, the duplicate tensor names should be resolved. Thanks.