Getting AttributeError: 'Tensor' object has no attribute 'log_prob' while saving a tensorflow model

Question

I am trying to save a variational autoencoder built using TensorFlow and Tensorflow probability. My goal is to serve the model using TensorFlow serving. I am open to serving the model in other fashion.

Here is my model:

def dense_layers(sizes):
return tfk.Sequential([tfkl.Dense(size, activation=tf.nn.leaky_relu) for size in sizes])

original_dim = 30
input_shape = (30,)
intermediary_dims = [20, 10, 8]
latent_dim = 2


prior = tfd.MultivariateNormalDiag(
        loc=tf.zeros([latent_dim]),
        scale_identity_multiplier=1.0)

encoder = tfk.Sequential([
    tfkl.InputLayer(input_shape=input_shape, name='encoder_input'),
    dense_layers(intermediary_dims),
    tfkl.Dense(tfpl.MultivariateNormalTriL.params_size(latent_dim), activation=None),
    tfpl.MultivariateNormalTriL(latent_dim, 
                           activity_regularizer=tfpl.KLDivergenceRegularizer(prior)),
], name='encoder')

encoder.summary()

decoder = tfk.Sequential([
    tfkl.InputLayer(input_shape=[latent_dim]),
    dense_layers(reversed(intermediary_dims)),
    tfkl.Dense(tfpl.IndependentNormal.params_size(original_dim), activation=None),
    tfpl.IndependentNormal(original_dim),
], name='decoder')

decoder.summary()

vae = tfk.Model(inputs=encoder.inputs,
                outputs=decoder(encoder.outputs[0]),
                name='vae_mlp')

negloglik = lambda x, rv_x: -rv_x.log_prob(x)

vae.compile(optimizer=tf.keras.optimizers.RMSprop(), 
            loss=negloglik)

vae.summary()

Here is how I am trying to save the model:

tf.keras.models.save_model(
                vae,
                "/opt/notebooks/saved/vae/1",
                overwrite=True,
                include_optimizer=True,
                save_format=None,
                signatures=None,
                options=None
            )

score 1 · Answer 1 · edited Oct 28 '20 at 21:49

The long term solution involves pushing for tensorflow bug #742 to be fixed. In the meantime, you can do what I did, modifying tensorflow code at file tensorflow_probability/python/layers/distribution_layer.py. What I did was:

Add a check in _make_kl_divergence_fn function, in order to avoid computing the divergence when the input is a Tensor object, like this:

    with tf.name_scope('kldivergence_loss'):
      if isinstance(distribution_a, tf.Tensor):
          return 0.0
      ...

Either change KLDivergenceRegularizer for a KLDivergenceAddLoss layer, or add a get_config method to the KLDivergenceRegularizer object:

  def get_config(self):
    config = {'use_exact_kl': self._use_exact_kl,
              'test_points_reduce_axis': self._test_points_reduce_axis,
              'weight': self._weight}
    return dict(list(config.items()))

Getting AttributeError: 'Tensor' object has no attribute 'log_prob' while saving a tensorflow model

1 Answers1