如何加载和使用预制的PyTorch InceptionV3模型对图像进行分类

时间:2018-12-19 04:57:50

标签: python pytorch torch

我遇到的问题与How can I load and use a PyTorch (.pth.tar) model相同,但没有被接受的答案,或者我可以弄清楚如何遵循给出的建议。

我是PyTorch的新手。我正在尝试加载此处引用的预训练PyTorch模型:https://github.com/macaodha/inat_comp_2018

我很确定我缺少一些胶水。

# load the model
import torch
model=torch.load("iNat_2018_InceptionV3.pth.tar",map_location='cpu')

# try to get it to classify an image
imsize = 256
loader = transforms.Compose([transforms.Scale(imsize), transforms.ToTensor()])

def image_loader(image_name):
    """load image, returns cuda tensor"""
    image = Image.open(image_name)
    image = loader(image).float()
    image = Variable(image, requires_grad=True)
    image = image.unsqueeze(0)  
    return image.cpu()  #assumes that you're using CPU

image = image_loader("test-image.jpg")

产生错误:

  

在()中   ----> 1个模型。预测(图像)

     

AttributeError:“ dict”对象没有属性“ predict”

1 个答案:

答案 0 :(得分:2)

问题

您的model实际上不是模型。保存后,它不仅包含参数,而且还包含与模型有关的其他信息,其形式类似于dict。

因此,torch.load("iNat_2018_InceptionV3.pth.tar")仅返回dict,当然它没有名为predict的属性。

model=torch.load("iNat_2018_InceptionV3.pth.tar",map_location='cpu')
type(model)
# dict

解决方案

在这种情况下(一般情况下),首先需要执行的是根据官方指南"Load models"实例化所需的模型类。

# First try
from torchvision.models import Inception3
v3 = Inception3()
v3.load_state_dict(model['state_dict']) # model that was imported in your code.

但是,直接输入model['state_dict']会引起一些有关Inception3参数形状不匹配的错误。

重要的是要知道Inception3在实例化之后被更改了什么。幸运的是,您可以在原始作者的train_inat.py中找到它。

# What the author has done
model = inception_v3(pretrained=True)
model.fc = nn.Linear(2048, args.num_classes) #where args.num_classes = 8142
model.aux_logits = False

现在我们知道要更改什么,让我们对第一次尝试进行一些修改。

# Second try
from torchvision.models import Inception3
v3 = Inception3()
v3.fc = nn.Linear(2048, 8142)
v3.aux_logits = False
v3.load_state_dict(model['state_dict']) # model that was imported in your code.

然后成功加载模型!

相关问题