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有解释:
因为BatchNorm
和Dropout
是随机的,会导致每次运行都不一致,所以在传递输入之前,需要切换至model.eval()
进行预测。或者使用model.train()
来进行训练。
网友评论