One to many mapping using probabilistic UNet

username: Assefa_Seyoum
cooked: <p>I have been working on this project for over a month now. However, I couldn’t make it work yet. So, I need help.</p>
<p>It is basically applying the <a href="https://arxiv.org/pdf/1806.05034.pdf" rel="noopener nofollow ugc">probabilistic UNet</a> paper on the leaf segmentation dataset.</p>
<p>Given an image, I want to output one instance at a time by sampling from a distribution.</p>
<p>The probabilistic UNet architecture is shown below.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/5/55fcb5135b8253cdc57c8c24becc4a5a4b719b2c.png" data-download-href="https://discuss.pytorch.org/uploads/default/55fcb5135b8253cdc57c8c24becc4a5a4b719b2c" title="prob unet"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/5/55fcb5135b8253cdc57c8c24becc4a5a4b719b2c_2_690x443.png" alt="prob unet" data-base62-sha1="cgG0uvX3qOiHDf5fekAQW92nil6" width="690" height="443" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/5/55fcb5135b8253cdc57c8c24becc4a5a4b719b2c_2_690x443.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/5/55fcb5135b8253cdc57c8c24becc4a5a4b719b2c_2_1035x664.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/5/55fcb5135b8253cdc57c8c24becc4a5a4b719b2c_2_1380x886.png 2x" data-dominant-color="EBEBEB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">prob unet</span><span class="informations">1761×1132 121 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>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.</p>
<p>What I want:</p>
<ol>
<li>Sample a noise</li>
<li>Get one instance<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/4/34111a8b0315628696a79b7d175c11c78cccd67d.png" data-download-href="https://discuss.pytorch.org/uploads/default/34111a8b0315628696a79b7d175c11c78cccd67d" title="whatIwant"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/4/34111a8b0315628696a79b7d175c11c78cccd67d_2_244x250.png" alt="whatIwant" data-base62-sha1="7qBtjY1Ou5A65LK8CyXqMHCrQap" width="244" height="250" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/4/34111a8b0315628696a79b7d175c11c78cccd67d_2_244x250.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/4/34111a8b0315628696a79b7d175c11c78cccd67d_2_366x375.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/3/4/34111a8b0315628696a79b7d175c11c78cccd67d_2_488x500.png 2x" data-dominant-color="EAE6EB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">whatIwant</span><span class="informations">1235×1265 118 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></li>
</ol>
<p>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 <a href="https://arxiv.org/pdf/1903.10145.pdf" rel="noopener nofollow ugc">beta scheduling</a> KL, and even removed the KL term, but the noise is still ignored.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/9/79d00889d289b5375538a7fe8eb188cb881bf788.png" data-download-href="https://discuss.pytorch.org/uploads/default/79d00889d289b5375538a7fe8eb188cb881bf788" title="whatIget"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/9/79d00889d289b5375538a7fe8eb188cb881bf788_2_225x249.png" alt="whatIget" data-base62-sha1="hnBtXUulFpN4EbyZvn8ZWMNq0u4" width="225" height="249" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/9/79d00889d289b5375538a7fe8eb188cb881bf788_2_225x249.png, https://discuss.pytorch.org/uploads/default/optimized/3X/7/9/79d00889d289b5375538a7fe8eb188cb881bf788_2_337x373.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/9/79d00889d289b5375538a7fe8eb188cb881bf788_2_450x498.png 2x" data-dominant-color="ECE9ED"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">whatIget</span><span class="informations">1257×1391 151 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Can someone be kind enough to review my code <a href="https://github.com/Assefa123/Probabilistic-UNet" rel="noopener nofollow ugc">here</a> or explain why this approach might not work?</p>
<p>Thank you.</p>

username: Jakub_Mitura
cooked: <p>Did you manage to make it work ?</p>

username: Assefa_Seyoum
cooked: <p>No, it didn’t work. Still don’t know why.</p>

username: zabboud
cooked: <p>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.</p>