CODEWITHSUNDEEP
pythonmachine-learningcomputer-visionpytorch
Mamen
Mamen

Reputation: 1436

Get probability from predicted class pytorch

I have a network like this for image classification, for now i have 2 classes:

class ActionNet(Module):

    def __init__(self, num_class=4):

        super(ActionNet, self).__init__()



        self.cnn_layer = Sequential(

            #conv1

            Conv2d(in_channels=1, out_channels=32, kernel_size=1, bias=False),

            BatchNorm2d(32),

            PReLU(num_parameters=32),

            MaxPool2d(kernel_size=3),

            #conv2

            Conv2d(in_channels=32, out_channels=64, kernel_size=1, bias=False),

            BatchNorm2d(64),

            PReLU(num_parameters=64),

            MaxPool2d(kernel_size=3),

            #flatten

            Flatten(),

            Linear(576, 128),

            BatchNorm1d(128),

            ReLU(inplace=True),

            Dropout(0.5),

            Linear(128, num_class)

        )



    def forward(self, x):

        x = self.cnn_layer(x)

        return x

then after i training my network, i predict the image using this code:

def predict_image(image):



    input = torch.from_numpy(image)

    input = input.unsqueeze(1)

    input = input.to(device)

    output = model(input)

    index = output.data.cpu().numpy().argmax()

    return index

how do i get all class probability of prediction image? so the result would be array of index with probaility like 0=0.1, 1=0.7

Upvotes: 2

Views: 4367

Answers (1)

Dishin H Goyani
Dishin H Goyani

Reputation: 7693

To get probability from model output here you can use softmax function.

Try this

import torch.nn.functional as F
...
prob = F.softmax(output, dim=1)
...

Upvotes: 1

Related Questions