pytorch indexing for pixel probabilities
Suppose I need to classify each pixel into one of 3 classes. I wish to get the probability of each pixel. Here is a minimal example. Question is, how do I get those probabilities.
import torch
import torch.nn.functional as F

y = torch.randint(0, 3, (2, 1, 5, 5)) # classes
logits = torch.randn(2, 3, 5, 5) 
prob  = F.softmax(logits, dim=1) # probability map

prob[y] # does not work


You are looking for torch.gather:
torch.gather(prob,1, y)

You gather the probabilities along the first dimension according to the indices of y.
