How to get the logit value in semanatic segmentation for the assigned class?

Hi everyone,

I got 4D logit value matrices with the shape (batchsize, classes (e.g. 9 logit values if number of classes ==9), height, width) and the assigned class matrices (label matrices) which are 3D and has the shape (batchsize, 1 (which are the indices of the true class), height, width). How can the logits values corresponding to the assigned class accessed easily. Basically, use the label matrices as index for the logit matrices. The label matrix additionally contains default values which are 255 to exclude these pixels from gradient calculation. Hence, these pixels has to be filtered out first.

Thank you in advance!


You could use the gather operation and filter out the specific values afterwards:

logits = torch.randn(2, 3, 4, 4)
target = torch.randint(0, 3, (2, 4, 4))
res = logits.gather(1, target.unsqueeze(1))
1 Like

Hi Peter,
thanks for your reply. I forgot to delete the question. I also used the gather operation and some masking to filter out values.

Best Matthias