U-Net output has 256 channels

Schroter: <p>Hi, I trained U-Net. But when given a test image, the output mask consists of 256 channels. Idealy it should be 1 right? The question I asked is below. Would anyone be able to help me in this matter please?<br>
Thanks &amp; Best Regards<br>
Schroter Michael</p>
<p>(<a href="https://stackoverflow.com/questions/73807983/predicted-mask-image-has-wrong-dimension-unet-typeerror-invalid-shape-2023-2" class="inline-onebox" rel="noopener nofollow ugc">python 3.x - Predicted mask image has wrong dimension unet- TypeError: Invalid shape (2023, 2023, 256) for image data - Stack Overflow</a>)</p>

ptrblck: <p>It depends on the use case as the output channels usually represent the number of classes in a multi-class segmentation use case. In your example you might be dealing with 256 classes.</p>

Schroter: <p>Hi, Thanks for the reply. In this case it is a cell segementation <a href="https://www.kaggle.com/competitions/hubmap-organ-segmentation" rel="noopener nofollow ugc">case</a>. So could it be different unconnected polygons is that also a posibility please?<br>
Thanks &amp; Best Regards<br>
Schroter Michael</p>

ptrblck: <p>I’m not familiar with this competition but by quickly skimming through the dataset it seems it’s a binary segmentation use case, so your number of output channels should be <code>1</code> with <code>nn.BCE(WithLogits)Loss</code> (this is the usual approach for a binary classification/segmentation) or <code>2</code> with <code>nn.CrossEntropyLoss</code> (this would be treating the use case as a 2-class multi-class segmentation).</p>

Schroter: <p>Hi, Thanks for the reply. <code>unet = UNet(in_channels=3, out_channels=3, init_features=8)</code> If use out_channels=1. I get an assertion error. However, my training loss also seems to be ok (less than .2). The output seems to be only the issue as far as I know please.<br>
Thanks &amp; Best Regards<br>
Schroter Michael</p>

ptrblck: <p>What kind of assertion error are you seeing? Could you post the error message and the stacktrace here, please?</p>

Schroter: <p>Hi, thanks for the reply. No worries. Please.</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
/tmp/ipykernel_17/1595321179.py in &lt;module&gt;
     25             with torch.set_grad_enabled(phase == "train"):
     26                 y_pred = unet(x)
---&gt; 27                 loss = dsc_loss(y_pred, y_true)
     28                 running_loss += loss.item()
     29 

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1108         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1109                 or _global_forward_hooks or _global_forward_pre_hooks):
-&gt; 1110             return forward_call(*input, **kwargs)
   1111         # Do not call functions when jit is used
   1112         full_backward_hooks, non_full_backward_hooks = [], []

/tmp/ipykernel_17/3304955502.py in forward(self, y_pred, y_true)
      7     def forward(self, y_pred, y_true):
      8         #print(f'x={x.shape}, pred={y_pred.shape}, y_true={y_true.shape}')
----&gt; 9         assert y_pred.size() == y_true.size()
     10         y_pred = y_pred[:, 0].contiguous().view(-1)
     11         y_true = y_true[:, 0].contiguous().view(-1)

AssertionError: 
</code></pre>
<p>Thanks &amp; Best Regards<br>
Schroter Michael</p>

ptrblck: <p><code>dsc_loss</code> seems to expect the same shapes for the model output and the target. I would guess it should be <code>[batch_size, 1, height, width]</code> so check the shape of both tensors and try to narrow down why they differ. Maybe a simple <code>unsqueeze</code> would be enough on one of the tensors.</p>

Schroter: <p>Hi, Thanks for the reply. I got this output which from a peice of code that I added just above <code>assert</code><br>
part. Here it is: <code>x=torch.Size([5, 3, 256, 256]), pred=torch.Size([5, 1, 256, 256]), y_true=torch.Size([5, 3, 256, 256])</code> Would this be of any help please.<br>
Thanks &amp; Best Regards<br>
Schroter Michael</p>

Schroter: <p>Hi, Thanks. I got the mask output with 1 dimension. But the array is filled with the value 1 please.<br>
Thanks &amp; Best Regards<br>
Schroter Michael</p>

ptrblck: <p>The mask seems to be an RGB mask? If so, then you would have to convert it to a binary mask by e.g. using thresholds.</p>

Schroter: <p>Hi,<br>
Thanks for the reply. The output channel was changed to 1. So the mask is one channel (2023,2023). Please refer the <a href="https://stackoverflow.com/questions/73817764/predicted-u-net-mask-is-all-filled-with-value-1" rel="noopener nofollow ugc">this link</a> for your kind understanding.<br>
Thanks &amp; Best Regards<br>
Schroter Michael</p>