Generating a model's confidence from a (1,2) pytorch tensor

Question

So in my classification neural network, my final tensor is

a = tensor([[ 546.3831, -796.4016]], grad_fn=<AddmmBackward>)

I would get which category the program thinks it is by calling

a.max(1)

But is there a way to calculate the confidence in which the network has in this decision? I am using this to generate a heat map type of thing, where the confidence of multiple images will come together to generate a heat map image.

Answer 1

You can use the softmax function for this:

from torch.functional import F

a = tensor([[ 546.3831, -796.4016]], grad_fn=<AddmmBackward>)
prob = F.softmax(a, dim=-1)

which outputs tensor([[1., 0.]]) .

Note that I'm assuming that your loss function is the cross-entropy loss (or similar) and that a is directly involved in the loss calculation.

If the above are true then your model was trained to use these numbers, a , to predict the probability of each class and my answer is totally valid, otherwise it is not.