我的电脑承担了它这个年纪不该承担的责任
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变量的值即为正确预测的数量。