python - Load a single image in a pretrained pytorch net

Question

Total newbie here, I'm using this pytorch SegNet implementation with a '.pth' file containing weights from a 50 epochs training. How can I load a single test image and see the net prediction? I know this may sound like a stupid question but I'm stuck. What I've got is:

from segnet import SegNet
import torch

model = SegNet(2)
model.load_state_dict(torch.load('./model_segnet_epoch50.pth'))

How do I "use" the net on a single test picture?

Answer 1

我提供了一个ResNet152预训练模型的示例。

def image_loader(loader, image_name):
    image = Image.open(image_name)
    image = loader(image).float()
    image = torch.tensor(image, requires_grad=True)
    image = image.unsqueeze(0)
    return image

data_transforms = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor()
])


model_ft = models.resnet152(pretrained=True)
model_ft.eval()

print( np.argmax(model_ft(image_loader(data_transforms, $FILENAME)).detach().numpy()))

$FILENAME是要加载的图像的路径和名称。我从这篇文章中得到了必要的帮助。

Answer 2

output = model(image) .

请注意，图像应该是一个Variable对象，并且输出也是如此。例如，如果您的图像是一个 Numpy 数组，您可以像这样转换它：

var_image = Variable(torch.Tensor(image))