Python 预测多张图片并将预测保存在阵列中

我目前正在寻找一种方法来预测多张图片，并将预测结果保存在一个数组中。我目前的方法是将所有图像张量保存在一个数组中，迭代数组，预测每个图像，并将输入保存在数组中。有没有更有效的方法

我的1图片预测代码

def predict_image(path):
    print("Prediction in progress")
    image = Image.open(path)

    transformation = transforms.Compose([
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    img_tensor = transformation(image).float()
    img_tensor = img_tensor.unsqueeze_(0)

    if torch.cuda.is_available():
        img_tensor.cuda()

    input = Variable(img_tensor)
    output = model(input)
    index = output.data.numpy().argmax()
    return index

---

我认为这可能是在测试时预测输出的最佳方法之一

您可能希望尝试的另一件事是以批处理方式提供输入。 由于将所有张量存储在一个列表中，因此可以使用

torch.stack（list）

将多个张量堆叠在一个唯一的张量中。然后你可以把它输入模型

注意，在这种情况下，你必须事先决定批量大小，你必须注意不要夸大。Pytork将张量存储在RAM（如果您使用的是CPU）或VRAM（如果您使用的是GPU）中，因此批量大小过大可能会导致CUDA OOM错误