pytorch 模型测试

发布于:2025-03-04 ⋅ 阅读:(16) ⋅ 点赞:(0)

在使用 PyTorch 进行模型测试时,一般包含加载测试数据、加载训练好的模型、进行推理以及评估模型性能等步骤。以下为你详细介绍每个步骤及对应的代码示例。

1. 导入必要的库

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

2. 加载测试数据

假设我们使用的是 CIFAR - 10 数据集作为示例,你需要定义数据预处理的转换操作,然后加载测试数据集。

# 定义数据预处理的转换操作
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 加载测试数据集
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=2)

# 类别标签
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

3. 定义模型结构

如果你已经有训练好的模型,这一步可以跳过。但为了完整性,这里给出一个简单的卷积神经网络(CNN)示例。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()

4. 加载训练好的模型

假设你已经将训练好的模型保存为 cifar_net.pth 文件,现在可以加载它。

# 加载模型
net.load_state_dict(torch.load('cifar_net.pth'))

5. 进行推理和评估

在测试阶段,我们需要将模型设置为评估模式,然后遍历测试数据集,对每个样本进行推理,并计算模型的准确率。

# 将模型设置为评估模式
net.eval()

correct =