在预训练的 pytorch 网络中加载单个图像

Posted 2023-03-12

【中文标题】在预训练的 pytorch 网络中加载单个图像

【英文标题】：Load a single image in a pretrained pytorch net

【发布时间】：2018-10-08 08:45:29

【问题描述】：

这里的总新手，我正在使用this pytorch SegNet 实现和一个“.pth”文件，其中包含来自 50 个 epoch 训练的权重。 如何加载单个测试图像并查看网络预测？ 我知道这听起来像是一个愚蠢的问题，但我被困住了。 我得到的是：

from segnet import SegNet
import torch

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

如何在单个测试图片上“使用”网络？

【参考方案1】：

output = model(image) .

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

var_image = Variable(torch.Tensor(image))

【讨论】：

注意：变量是deprecated in PyTorch 0.4。

【参考方案2】：

我提供了一个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 是要加载的图像的路径和名称。我从这个post 得到了必要的帮助。

【讨论】：

您如何使用自己编写的模型来做到这一点？

@joehoeller 不确定您的确切意思。你不能把你的模型分配给model_ft吗？