One to many mapping using probabilistic UNet

I have been working on this project for over a month now. However, I couldn’t make it work yet. So, I need help.

It is basically applying the probabilistic UNet paper on the leaf segmentation dataset.

Given an image, I want to output one instance at a time by sampling from a distribution.

The probabilistic UNet architecture is shown below.

The architecture has two loss components; the KL divergence and reconstruction loss. When I train it, the KL loss goes to zero, however the reconstruction loss works differently that what I want.

What I want:

  1. Sample a noise
  2. Get one instance

What actually happens is that the noise is ignored it acts like a semantic segmentation. Basically, the RGB image and the different instances are mapped to the same latent space. I have tried beta scheduling KL, and even removed the KL term, but the noise is still ignored.

Can someone be kind enough to review my code here or explain why this approach might not work?

Thank you.

1 Like

Did you manage to make it work ?

No, it didn’t work. Still don’t know why.

1 Like

I’m not sure if this is still relevant, but I had a similar issue, and the only way to resolve it was to significantly increase the Beta term, basically forcing the posterior to the prior – by increasing it I mean to a factor of 10. the original paper mentions a Beta of 10, but I only managed to get results after increasing the Beta to 100 - I haven’t attempted tuning to find best performance yet, but this should be a good starting point for you.