尝试预测单个图像 CNN pytorch 时出错

问题描述

错误信息

回溯（最近一次调用最后一次）： 文件“pred.py”，第 134 行，在 输出 = 模型（数据） 运行时错误：4 维权重 [16,3,3] 的预期 4 维输入，但得到大小为 [1,32,32] 的 3 维输入。

预测代码

normalize = transforms.normalize(mean=[0.4914,0.4824,0.4467],std=[0.2471,0.2435,0.2616])
train_set = transforms.Compose([
                                 transforms.RandomCrop(32,padding=4),transforms.RandomHorizontalFlip(),transforms.ToTensor(),normalize,])

model = models.condensenet(args)
model = nn.DataParallel(model)
PATH = "results/savedir/save_models/checkpoint_001.pth.tar"

model.load_state_dict(torch.load(PATH)['state_dict'])


device = torch.device("cpu")

model.eval()

image = Image.open("horse.jpg")
input = train_set(image)
train_loader = torch.utils.data.DataLoader(
        input,batch_size=1,shuffle=True,num_workers=1)
for i,data in enumerate(train_loader):
    
    #input_var = torch.autograd.Variable(data,volatile=True)
    #input_var = input_var.view(1,32)
    
    **output = model(data)
topk=(1,5)
maxk = max(topk)

_,pred = output.topk(maxk,1,True,True)
    

Am getting this error when am trying to predict on a single image Image shape/size error message

Link to saved model

Training code repository

解决方法

请取消注释此行 #input_var = input_var.view(1,3,32,32) 以便您的输入维度为 4。
我假设你没有。输入通道数为 3，如果是 1，则使用 input_var = input_var.view(1,1,32) 如果灰度

,

没有执行 for 循环和 train_loader，而是通过将输入直接传递到模型中来解决此问题。像这样

input = train_set(image)
input = input.unsqueeze(0)
model.eval()
output = model(input)

可以在此处找到更多详细信息link