Bayesian CNN keeps giving the same output regardless of the input

[Chaudry24]

Hi, guys!

I am trying to use tensorflow probability for some Bayesian prediction but there is a problem. I have trained three different models so far, and each of these end up giving me the same output regardless of the input.

The three models are as follows:

• model 1: 3 layers of Convolution2dReparametrization with custom kernels.
• model 2: 3 layers of Convolution2dReparametrization with default kernels,
• model 3: 1 layer of Convolution2dReparametrization and 2 layers of Conv2d.

My code is here: https://pastebin.com/wtYuaB6j

I am only presenting three models here but I have tried even more than 3 and the result is still the same. As soon as I incorporate even a single tfp layer, eventually my model starts predicting the same output. I have tried changing the optimizer, the loss, the learning_rate and nothing has changed so far.

tensorflow version is 2.7.0 tensorflow_probability version is 0.15.0

Any help or suggestions are appreciated. If all else fails, this might be a bug in tfp.

[ColCarroll]

Thanks for providing some code! Can you point out which parts of the pastebin you're running, what the output looked like, and what you were expecting? I'm not quite sure what "eventually my model starts predicting the same output" means!

[Chaudry24]

Hi, @ColCarroll! Thank you for your response.

I am running the whole thing in google colab notebook. Here is what I mean by predicting the same output, my input data is a N x 16 x 16 x 1 where N is batch size. My output is supposed to be a 3 x 1 array. When I train the model, I expect to see different numbers in output, but I end up seeing, say [1, 1, 1] regardless of what I input into the model.

I think I may have a possible solution to this. It has to do with the likelihood function. Since the loss I am equipping my model with is MeanAbsoluteError(), I think there is something weird going on because VAE usually uses a maximum likelihood loss intead