Decoder not learning

Yazan_Maarouf: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/3/239d6dd446848da6aef2c4deca417c6024939442.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/239d6dd446848da6aef2c4deca417c6024939442" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/239d6dd446848da6aef2c4deca417c6024939442_2_572x500.jpeg" alt="image" data-base62-sha1="5540661V9IPBvYKxpeICVGPlqGC" width="572" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/239d6dd446848da6aef2c4deca417c6024939442_2_572x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/239d6dd446848da6aef2c4deca417c6024939442_2_858x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/239d6dd446848da6aef2c4deca417c6024939442_2_1144x1000.jpeg 2x" data-dominant-color="EAEBEF"><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">image</span><span class="informations">1920×1677 137 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><br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/6/b/6b92275eb91df7520150d3a2f639dc79669b9764.png" alt="image" data-base62-sha1="flCdMR3XzBvy5ens5rH4M7gmfnC" width="628" height="320"></p>
<p>I am working on a model where the task is panoptic segmentation (I think thats what its called).<br>
The <code>input</code> is an image and the <code>output</code> is a <strong>3D Matrix</strong> where each 2D matrix is a class (I have 13 classes). I am trying to overfit a single image</p>
<p>The model is as follows:</p>
<ol>
<li>Encoder (MobileNetV4)</li>
<li>Decoder (MobileNetV4 reversed with ConvTranspose2D with skip connections)</li>
<li>Transformer (Takes in encoder output and outputs weights → dynamic weights)</li>
<li>DynamicConv (Layer that takes the Decoder output and the Transformer’s outputs)<br>
The Dynamic Conv layer uses the transformer outputs as weights over the Decoder output.<br>
Decoder output might look like 1,16,256,256<br>
Transformer output looks like this: 1, 16, 13<br>
The final output is 1, 13, 256, 256.<br>
The screenshot above shows the decoder layer always has 0 for the gradients. Mind you, this is epoch 1000. But Ive also plotted the 3D gradients in earlier layers. The Transformer had bigger gradients as well. However, the Decoder layer always had 0 since the beginning. So the Transformer learned a bit?, but the Decoder did not learn.<br>
I know my issue is related to the Decoder weights and how we are updating the weights for the decoder, I just need a sanity check from someone. Im not a researcher and just constantly trying to learn.</li>
</ol>
<p>Early epoch:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/6/26eb888c830aa072f2ef324d4c371f4e4dcd1200.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/26eb888c830aa072f2ef324d4c371f4e4dcd1200" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/6/26eb888c830aa072f2ef324d4c371f4e4dcd1200_2_544x500.jpeg" alt="image" data-base62-sha1="5yiLQYCTIDxJ7JoMiKEVrXn5H4Q" width="544" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/6/26eb888c830aa072f2ef324d4c371f4e4dcd1200_2_544x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/2/6/26eb888c830aa072f2ef324d4c371f4e4dcd1200_2_816x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/6/26eb888c830aa072f2ef324d4c371f4e4dcd1200_2_1088x1000.jpeg 2x" data-dominant-color="E2E4EB"><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">image</span><span class="informations">1864×1712 203 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>Am I right to assume that the model is not using the UNET Decoder and the Transformer  effectively when learning ?</p>
<p>This is 1000 epoch to fit a single image with its respective 3D masks matrix</p>

Yazan_Maarouf: <p>I found the issues in my code:</p>
<ol>
<li>
<p>I was not applying convolutions correctly for my DynamicConvolution → had to use conv groups to apply the weights on the output</p>
</li>
<li>
<p>My standardization was applied both on my input and output (output was already scaled)</p>
</li>
</ol>
<p>I shouldve had more doubt about my overfitting. It took a long time to overfit 1 example → that’s not good at all</p>