PyTorch计算机视觉入门：测试模型与评估，对单帧图片进行推理

在完成模型的训练之后，对模型进行测试与评估是至关重要的一步，它能帮助我们理解模型在未知数据上的泛化能力。本篇指南将带您了解如何使用PyTorch进行模型测试，并对测试结果进行分析。我们将基于之前训练好的模型，演示如何加载数据、进行预测、计算指标以及可视化结果。

准备工作

假设您已经有一个训练好的模型，保存在.pth文件中，以及一个用于测试的自定义数据集。我们将继续使用前文提到的自定义数据集CustomDataset类，并引入一些新的概念和代码。

加载测试数据集

与训练过程类似，首先需要加载测试数据集，并对其进行适当的预处理。确保您的测试集遵循与训练集相同的数据结构和预处理步骤。

test_dataset = CustomImageDataset(data_path="./data/", model= "test", transform = transform)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=True)

测试模型

训练完成后，使用测试数据集来评估模型的性能

def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target)
            # max 函数返回两个值，一个是是数值，一个是index
            pred = output.max(1, keepdim=True)[1] # 找到概率最大的下标 

            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

Test set: Average loss: 0.0018, Accuracy: 9671/10000 (97%)

单帧图片进行测试

# test single image 
img = Image.open("./data/data_test/1.jpg")
img_t = transform(img)
img_t = img_t.unsqueeze(0)  # 变为[1, 1, 28, 28]
img_t = img_t.to(device)
model.eval()
output = model(img_t)
_, predicted_class = torch.max(output, 1)
print(predicted_class)

tensor([2], device='cuda:0')

通过以上步骤，我们可以全面地评估和分析PyTorch模型在计算机视觉任务中的表现，从而确保模型在实际应用中的有效性和可靠性。

关注我的公众号Ai fighting, 第一时间获取更新内容。