U-Net output has 256 channels

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?
Thanks & Best Regards
Schroter Michael

(python 3.x - Predicted mask image has wrong dimension unet- TypeError: Invalid shape (2023, 2023, 256) for image data - Stack Overflow)

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.

Hi, Thanks for the reply. In this case it is a cell segementation case. So could it be different unconnected polygons is that also a posibility please?
Thanks & Best Regards
Schroter Michael

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 1 with nn.BCE(WithLogits)Loss (this is the usual approach for a binary classification/segmentation) or 2 with nn.CrossEntropyLoss (this would be treating the use case as a 2-class multi-class segmentation).

Hi, Thanks for the reply. unet = UNet(in_channels=3, out_channels=3, init_features=8) 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.
Thanks & Best Regards
Schroter Michael

What kind of assertion error are you seeing? Could you post the error message and the stacktrace here, please?

Hi, thanks for the reply. No worries. Please.

AssertionError                            Traceback (most recent call last)
/tmp/ipykernel_17/1595321179.py in <module>
     25             with torch.set_grad_enabled(phase == "train"):
     26                 y_pred = unet(x)
---> 27                 loss = dsc_loss(y_pred, y_true)
     28                 running_loss += loss.item()

/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):
-> 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}')
----> 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)


Thanks & Best Regards
Schroter Michael

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

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

Hi, Thanks. I got the mask output with 1 dimension. But the array is filled with the value 1 please.
Thanks & Best Regards
Schroter Michael

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.

Thanks for the reply. The output channel was changed to 1. So the mask is one channel (2023,2023). Please refer the this link for your kind understanding.
Thanks & Best Regards
Schroter Michael