为什么我的预测函数会出错？ ValueError：没有足够的值来解包（预期 2，得到 1）

Question

我正在尝试使用预训练模型进行预测，该模型使用 UNET 和 pytorch 进行二进制分割。这是我的代码： model.eval() # 设置模型为评估模式

class SimDataset(Dataset):
    def __init__(self, path, transform=None, isMask=False):
                
        self.m = ("test")
        self.path = path

        self.transform = transform
        self.isMask = isMask
    
    def __len__(self):
        return len(self.path)
    
    def __getitem__(self, idx):     
        one_image = os.path.join(self.m, self.path[idx]) # preparing image path/location
        img_temp = Image.open(one_image) # load RGB input image
        

        if self.transform:
            image = self.transform(img_temp)
      

        input_image = np.array(img_temp).astype('float32') # converting one image to np array
        input_image = np.transpose(input_image, (2, 0 ,1)) # converting from hwc to chw [(256,256,3) => (3, 256, 256)]
                
        return [input_image]

testlist = list(os.listdir(r"test"))
len(testlist)
image_datasets = {
    'testlist': testlist
}

dataset_sizes = {
    x: len(image_datasets[x]) for x in image_datasets.keys()
}
test_dataset = SimDataset(testlist, transform = trans, isMask=False)
test_loader = DataLoader(test_dataset, batch_size=3, shuffle=False, num_workers=0)
inputs, labels = next(iter(test_loader))
inputs = inputs.to(device)
labels = labels.to(device)

pred = model(inputs)
pred = torch.sigmoid(pred)
pred = pred.data.cpu().numpy()
print(pred.shape)

但它显示错误提示 -> ValueError: not enough values to unpack (expected 2, got 1).

Answer 1

您的代码需要来自数据加载器的两个输出：

inputs, labels = next(iter(test_loader))

但是，您的数据集中的__getitem__ 方法仅返回一个单个 输出：

return [input_image]

您要么从__getitem__ 返回两个 输出，包括图像和标签。 或者只期望test_loader 的单个输出。