Body: <p>I’m training a U-Net (model below) to predict 4 heatmaps (gaussian centered around a keypoint, one in each channel). Each channel is for some reason outputting the same result, an example is given of a test image where the blue is ground truth for that channel and red is the output of the u-net.
  I have tried using L1, MSE and adaptive wing loss (Wang 2019), and none are able to regress the heatmaps.
  I'm not sure what I'm doing wrong would appreciate any advice. Thanks</p>
  <p><a href="https://i.sstatic.net/yBtU7.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/yBtU7.png" alt="" /></a></p>
  <pre><code>class CNN(nn.Module):
      def __init__(self):
          super(CNN,self).__init__()
          self.layer1 = nn.Sequential(
              nn.Conv2d(config.numinputs, 64,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(64))
      
          self.layer2 = nn.Sequential( 
              nn.Conv2d(64, 64,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(64))
          
          self.layer3 = nn.Sequential(
              nn.MaxPool2d(2, stride=2, padding=0))
              
          self.layer4 = nn.Sequential(
              nn.Conv2d(64,128,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(128))
  
          self.layer5 = nn.Sequential(
              nn.Conv2d(128, 128,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(128))
              
          self.layer6 = nn.Sequential(
              nn.MaxPool2d(2, stride=2, padding=0))
              
          self.layer7 = nn.Sequential(
              nn.Conv2d(128, 256,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(256))
              
          self.layer8 = nn.Sequential(
              nn.Conv2d(256, 256,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(256))
              
          self.layer9 = nn.Sequential(
              nn.Conv2d(256, 256,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(256))
              
          self.layer10 = nn.UpsamplingBilinear2d(scale_factor=2)
              
          self.layer11 = nn.Sequential(
              nn.Conv2d(256, 128,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(128))
              
          self.layer12 = nn.Sequential(
              nn.Conv2d(128, 128,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(128))
              
          self.layer13 = nn.Sequential(
              nn.Conv2d(128, 128,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(128))
              
          self.layer14 = nn.UpsamplingBilinear2d(scale_factor=2)
              
          self.layer15 = nn.Sequential(
              nn.Conv2d(128, 64,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(64))
              
          self.layer16 = nn.Sequential(
              nn.Conv2d(64, 64,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(64))
              
          self.layer17 = nn.Sequential(
              nn.Conv2d(64, 4,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(1))
          
          self.layer18 = nn.Softmax(dim=1)
              
              
      def forward(self,x):
          x = self.layer1(x)  
          x = self.layer2(x)  
          x = self.layer3(x)  
          x = self.layer4(x)
          x = self.layer5(x)
          x = self.layer6(x)  
          x = self.layer7(x)
          x = self.layer8(x)
          x = self.layer9(x)  
          x = self.layer10(x)  
          x = self.layer11(x)  
          x = self.layer12(x)  
          x = self.layer13(x)  
          x = self.layer14(x)  
          x = self.layer15(x)  
          x = self.layer16(x)
          x = self.layer17(x) 
          x = x.view(1, -1)
          x = self.layer18(x)
          x = x.reshape(4, 512, 512)
          
          return x
  
  </code></pre>
  <hr />
  <p>Other test images outputted
  ...</p>
  <p><a href="https://i.sstatic.net/yBtU7.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/yBtU7.png" alt="" /></a></p>
  <p><a href="https://i.sstatic.net/ba1Yh.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/ba1Yh.png" alt="" /></a></p>
  <p><a href="https://i.sstatic.net/6JKoN.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/6JKoN.png" alt="" /></a></p>
  
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2022-04-16T15:16:16.250Z
Comments.Text: It seems like it gets the first channel right. Is this the case for other samples as well? Also, should we assume `config.numinputs` is equal to 3?
Comments.UserDisplayName: Andrei-Cristian Rad
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2022-04-18T10:16:53.117Z
Comments.Text: @Andrei-CristianRad actually no - config.numinputs is only 1 (the input frame to the CNN). I've put another comment in this thread to show the other output images. ground truth is a gaussian around the blue points
Comments.UserDisplayName: Solace
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2022-11-09T23:19:00.087Z
Comments.Text: Did you ever figure out the problem? I'm working on something similar... I suspect I know the problem which is why I'm asking, but I don't know the solution. Basically with landmark prediction it's like 1 or a few pixels with a certain tag, and the rest are off. In your case with 500x500 images that's a few pixels out of 250k. So the clear local minima is to have them all turned off for all the tags. I'm guessing you need to either use some sort of different loss, or proper re-scaling of your loss. But not sure about which loss, and the rescaling seems ad-hoc.
Comments.UserDisplayName: Sergey Bokhnyak
Tags: <convolutional-neural-networks><pytorch><u-net>
Title: Training a u-net for multi-landmark heatmap regression producing the same heatmap for each channel
UserDisplayName: Solace
ThreadText: # Training a u-net for multi-landmark heatmap regression producing the same heatmap for each channel
  
  (Asked by: username_0 on 2022-04-15)
  
  I’m training a U-Net (model below) to predict 4 heatmaps (gaussian centered around a keypoint, one in each channel). Each channel is for some reason outputting the same result, an example is given of a test image where the blue is ground truth for that channel and red is the output of the u-net.
  I have tried using L1, MSE and adaptive wing loss (Wang 2019), and none are able to regress the heatmaps.
  I'm not sure what I'm doing wrong would appreciate any advice. Thanks
  
  [![](https://i.sstatic.net/yBtU7.png)](https://i.sstatic.net/yBtU7.png)
  
  ```
  class CNN(nn.Module):
      def __init__(self):
          super(CNN,self).__init__()
          self.layer1 = nn.Sequential(
              nn.Conv2d(config.numinputs, 64,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(64))
      
          self.layer2 = nn.Sequential( 
              nn.Conv2d(64, 64,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(64))
          
          self.layer3 = nn.Sequential(
              nn.MaxPool2d(2, stride=2, padding=0))
              
          self.layer4 = nn.Sequential(
              nn.Conv2d(64,128,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(128))
  
          self.layer5 = nn.Sequential(
              nn.Conv2d(128, 128,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(128))
              
          self.layer6 = nn.Sequential(
              nn.MaxPool2d(2, stride=2, padding=0))
              
          self.layer7 = nn.Sequential(
              nn.Conv2d(128, 256,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(256))
              
          self.layer8 = nn.Sequential(
              nn.Conv2d(256, 256,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(256))
              
          self.layer9 = nn.Sequential(
              nn.Conv2d(256, 256,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(256))
              
          self.layer10 = nn.UpsamplingBilinear2d(scale_factor=2)
              
          self.layer11 = nn.Sequential(
              nn.Conv2d(256, 128,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(128))
              
          self.layer12 = nn.Sequential(
              nn.Conv2d(128, 128,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(128))
              
          self.layer13 = nn.Sequential(
              nn.Conv2d(128, 128,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(128))
              
          self.layer14 = nn.UpsamplingBilinear2d(scale_factor=2)
              
          self.layer15 = nn.Sequential(
              nn.Conv2d(128, 64,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(64))
              
          self.layer16 = nn.Sequential(
              nn.Conv2d(64, 64,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(64))
              
          self.layer17 = nn.Sequential(
              nn.Conv2d(64, 4,kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
              nn.ReLU(),
              nn.BatchNorm2d(1))
          
          self.layer18 = nn.Softmax(dim=1)
              
              
      def forward(self,x):
          x = self.layer1(x)  
          x = self.layer2(x)  
          x = self.layer3(x)  
          x = self.layer4(x)
          x = self.layer5(x)
          x = self.layer6(x)  
          x = self.layer7(x)
          x = self.layer8(x)
          x = self.layer9(x)  
          x = self.layer10(x)  
          x = self.layer11(x)  
          x = self.layer12(x)  
          x = self.layer13(x)  
          x = self.layer14(x)  
          x = self.layer15(x)  
          x = self.layer16(x)
          x = self.layer17(x) 
          x = x.view(1, -1)
          x = self.layer18(x)
          x = x.reshape(4, 512, 512)
          
          return x
  
  ```
  
  ---
  
  Other test images outputted
  ...
  
  [![](https://i.sstatic.net/yBtU7.png)](https://i.sstatic.net/yBtU7.png)
  
  [![](https://i.sstatic.net/ba1Yh.png)](https://i.sstatic.net/ba1Yh.png)
  
  [![](https://i.sstatic.net/6JKoN.png)](https://i.sstatic.net/6JKoN.png)
  
  **Tags:** convolutional-neural-networks, pytorch, u-net
  
  > **Comments:**
  > **username_1**: It seems like it gets the first channel right. Is this the case for other samples as well? Also, should we assume `config.numinputs` is equal to 3?  
  > **username_0**: @username_1 actually no - config.numinputs is only 1 (the input frame to the CNN). I've put another comment in this thread to show the other output images. ground truth is a gaussian around the blue points  
  > **username_2**: Did you ever figure out the problem? I'm working on something similar... I suspect I know the problem which is why I'm asking, but I don't know the solution. Basically with landmark prediction it's like 1 or a few pixels with a certain tag, and the rest are off. In your case with 500x500 images that's a few pixels out of 250k. So the clear local minima is to have them all turned off for all the tags. I'm guessing you need to either use some sort of different loss, or proper re-scaling of your loss. But not sure about which loss, and the rescaling seems ad-hoc.  
  
  ---
  
  ## 0 answers
  
  