用cifar-10数据测试VGG-16的准确性

我的电脑承担了它这个年纪不该承担的责任

import torch
import torchvision
import torchvision.transforms as transforms
from torchvision.models import vgg16

# 加载CIFAR-10数据集并进行预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 设置较小的数据集大小
subset_indices = range(500)  # 选择前500个示例作为子集
subset_sampler = torch.utils.data.SubsetRandomSampler(subset_indices)

testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
batch_size = 16  # 设置较小的批次大小

testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, sampler=subset_sampler, shuffle=False, num_workers=0)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 加载预训练的VGG-16模型
model = vgg16(pretrained=True)

# 修改最后一层全连接层的输出维度
model.classifier[6] = torch.nn.Linear(4096, 10)

model.to(device)
model.eval()  # 切换为评估模式

correct = 0
total = 0

with torch.no_grad():
    for data in testloader:
        images, labels = data[0].to(device), data[1].to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print('Accuracy of the VGG-16 model on the CIFAR-10 test images: %.2f %%' % accuracy)

 

首先，我们使用torch.max()函数找到输出张量outputs中每个样本预测的最大值及其对应的索引。预测的结果保存在变量predicted中。

然后，我们将预测结果predicted与真实标签labels进行比较。(predicted == labels)会产生一个布尔张量，其中相等的位置为True，不相等的位置为False。

最后，使用.sum().item()计算True的数量，即正确预测的数量，并将其累加到correct变量中。

在循环结束后，correct变量的值即为正确预测的数量。