我有一个用于图像分类的网络，目前有两个类别：

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

训练网络后，我使用以下代码预测图像：

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

如何获取预测图像的所有类别概率？结果应为带概率的索引数组，如 0=0.1, 1=0.7

回答：

要从模型输出中获取概率，可以使用 softmax 函数。

尝试如下代码：

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