sacmehta / ESPNetv2

A light-weight, power efficient, and general purpose convolutional neural network
MIT License
447 stars 70 forks source link

Is it possible to get probability(likelihood) for foreground class label? #18

Closed c13proto closed 5 years ago

c13proto commented 5 years ago

Hello Sachin, I am very beginner of Deep Learning , I apologize if my question is stupid.

I trained with my own annotation images (background is 0, foreground is 1 ) and set "--classes=2". Then I want to know the probability(likelihood) of label 1 for each pixel of test images. In other word, I want to get heat-map of probability of foreground label. Is it possible? And, I want to know whether trying 2 class segmentation on this network is correct approach or not.

sacmehta commented 5 years ago

For computing probability, first apply softmax in output tensor and then take torch.max of the resultant tensor. This will return two values, max probability along the channel and corresponding index value (or class).

See PyTorch documentation for more details

c13proto commented 5 years ago

I see, thank you for your answer. I try to read it!