How to save a model with DenseVariational layer?

Viewed 298

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?

2 Answers

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

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.

Related