How to correctly use DenseFlipout layers with TensorFlow Probability-CodePudding

i am a novice with both TensorFlow and TensorFlow Probability. I am using this network for a regression task.

def normal_sp(params): 
  return tfd.Normal(loc=params[:,0:1], scale=1e-3   tf.math.softplus(0.05 * params[:,1:2]))

kernel_divergence_fn=lambda q, p, _: tfp.distributions.kl_divergence(q, p) / (x.shape[0] * 1.0)
bias_divergence_fn=lambda q, p, _: tfp.distributions.kl_divergence(q, p) / (x.shape[0] * 1.0)


inputs = Input(shape=(1,),name="input layer")


hidden = tfp.layers.DenseFlipout(50,bias_posterior_fn=tfp.layers.util.default_mean_field_normal_fn(),
                           bias_prior_fn=tfp.layers.default_multivariate_normal_fn,
                           kernel_divergence_fn=kernel_divergence_fn,
                           bias_divergence_fn=bias_divergence_fn,activation="relu",name="DenseFlipout_layer_1")(inputs)
hidden = tfp.layers.DenseFlipout(100,bias_posterior_fn=tfp.layers.util.default_mean_field_normal_fn(),
                           bias_prior_fn=tfp.layers.default_multivariate_normal_fn,
                           kernel_divergence_fn=kernel_divergence_fn,
                           bias_divergence_fn=bias_divergence_fn,activation="relu",name="DenseFlipout_layer_2")(hidden)
hidden = tfp.layers.DenseFlipout(100,bias_posterior_fn=tfp.layers.util.default_mean_field_normal_fn(),
                           bias_prior_fn=tfp.layers.default_multivariate_normal_fn,
                           kernel_divergence_fn=kernel_divergence_fn,
                           bias_divergence_fn=bias_divergence_fn,activation="relu",name="DenseFlipout_layer_3")(hidden)
params = tfp.layers.DenseFlipout(2,bias_posterior_fn=tfp.layers.util.default_mean_field_normal_fn(),
                           bias_prior_fn=tfp.layers.default_multivariate_normal_fn,
                           kernel_divergence_fn=kernel_divergence_fn,
                           bias_divergence_fn=bias_divergence_fn,name="DenseFlipout_layer_4")(hidden)
dist = tfp.layers.DistributionLambda(normal_sp)(params) 


model_vi = Model(inputs=inputs, outputs=dist)
model_vi.compile(Adam(learning_rate=0.002), loss=NLL) 

model_params = Model(inputs=inputs, outputs=params)

my question is related to the loss function:

in the example posted here, the authors add the kl divergence to the loss function https://www.tensorflow.org/probability/api_docs/python/tfp/layers/DenseFlipout

kl = sum(model.losses)
loss = neg_log_likelihood   kl

but in the example here https://colab.research.google.com/github/tensorchiefs/dl_book/blob/master/chapter_08/nb_ch08_03.ipynb

the loss function is simply the NLL. My question is : do i have to add manually the kl divergence or does tensorflow calculate it automatically? in the first case, how do i do it since model.losses doesn't seem to work? Thanks to anyone who help

CodePudding user response：

If you're using Keras to train, the per-layer losses (KLs) are included in the overall loss (I am 90% sure this is right -- you could check by overriding the kl_divergence_fn to return some absurd value and see if your overall loss becomes absurd).

In the example from the docs (which are, ahem, a bit ancient), keras is not doing the training; instead an optimizer is being applied to a manually written loss, and so one has to grab all the per layer losses and add them in.