One to many mapping using probabilistic UNet

vision
1

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.

prob unet
prob unet1761×1132 121 KB

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
    whatIwant
    whatIwant1235×1265 118 KB

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.

whatIget
whatIget1257×1391 151 KB

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

Thank you.

1 Like
2

Did you manage to make it work ?

3

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

1 Like
4

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.