我有一个这样的图像分类网络,目前我有两个类:
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的索引
要从模型输出中获取概率,可以使用^{} 函数
试试这个
相关问题 更多 >
编程相关推荐