美文网首页
Pytorch使用预训练模型时,同一张图片预测不同的结果或者很多

Pytorch使用预训练模型时,同一张图片预测不同的结果或者很多

作者: a326bb0fe451 | 来源:发表于2018-08-30 16:12 被阅读480次
from torchvision import datasets, models, transforms

alexnet_model = models.alexnet(pretrained=True)
alexnet_model.eval()

transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
])

img = io.imread("cat.jpg")
img = transform(img)
img = img.unsqueeze(0)

prob = F.softmax(alexnet_model(image), dim=1)
pred = torch.argmax(prob, dim=1)

以上代码使用pytorch预训练的alexnet进行预测一张猫的图片,如果不加上alexnet_model.eval()则会导致每次运行预测的结果都不一样。
具体原因在Pretrained resnet constant output有解释:

因为BatchNormDropout是随机的,会导致每次运行都不一致,所以在传递输入之前,需要切换至model.eval()进行预测。或者使用model.train()来进行训练。

相关文章

网友评论

      本文标题:Pytorch使用预训练模型时,同一张图片预测不同的结果或者很多

      本文链接:https://www.haomeiwen.com/subject/ktkdwftx.html