How to save the predicted multi-classes segmentation outputs tensor like the single channel ground truth?

vision
#1

In multi-classes segmentation, the shape of the model’s last layer is [batch , num_classes, H, W] , e.g. [1, 6, 512, 512],however ,the ground truth shape is [1, 1, 512, 512], how to transform the output tensor shape
into the ground truth shape and save it?
Thanks in advance!

#2

If you are dealing with a multi-class classification and would like to use nn.CrossEntropyLoss (or nn.NLLLoss) as your criterion, the output of your model should contain the logits (or log probabilities) as [batch_size, nb_classes, height, width], while your target should be a LongTensor containing the class indices as [batch_size, height, width].
If your target already contains these indices, you can just call target = target.squeeze(1) on it to remove the channel dimension.

#3

I have known how to do, maybe I have no clear statement, you misunderstand my question , we should use torch.argmax() on the predicted multi-classes segmentation outputs

1 Like
#4

Hi, I am in the same situation. Finally you used softmax/argmax in the num_classes dimension? [batch , num_classes, H, W -> [batch , H, W]