PyTorch 模型测试和性能评估代码示例

import torch
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from mydataset import MyDataset
from torch.utils.data import DataLoader
from sklearn.metrics import accuracy_score, confusion_matrix
import argparse
import torch.nn as nn


def test(model, dataset, criterion, num_classes):
    model.eval()
    total_batch_num = 0.
    val_loss = 0
    prediction = []
    labels = []
    feature_list = torch.tensor([])
    if torch.cuda.is_available():
        feature_list = feature_list.cuda()
    class_correct = [0] * num_classes  # 初始化每个类别的正确预测数为0
    class_total = [0] * num_classes  # 初始化每个类别的样本总数为0
    for (step, i) in enumerate(dataset):
        total_batch_num += 1
        batch_x = i['data']
        batch_y = i['label'] - 1
        batch_x = torch.unsqueeze(batch_x, dim=1)
        batch_x = batch_x.float()
        if torch.cuda.is_available():
            batch_x = batch_x.cuda()
            batch_y = batch_y.cuda()
        feature, probs = model(batch_x)
        batch_label = batch_y.unsqueeze(1).float()
        feature_label = torch.cat((feature, batch_label), dim=1)
        feature_list = torch.cat((feature_list, feature_label), dim=0)
        loss = criterion(probs, batch_y)
        _, pred = torch.max(probs, dim=1)
        predi = pred.tolist()
        label = batch_y.tolist()
        val_loss += loss.item()
        prediction.extend(predi)
        labels.extend(label)
        # 计算每个类别的正确预测数和样本总数
        for j in range(len(label)):
            if predi[j] == label[j]:
                class_correct[label[j]] += 1
            class_total[label[j]] += 1
    accuracy = accuracy_score(labels, prediction)
    C = confusion_matrix(labels, prediction)
    class_accuracy = [class_correct[i] / class_total[i] for i in range(num_classes)]  # 计算每个类别的精确率
    return accuracy, val_loss / total_batch_num, feature_list, C, class_accuracy

# 调用测试函数
parser = argparse.ArgumentParser()
rootpath = 'das_data'
parser.add_argument('--root2', type=str, default=rootpath + '/test', help='rootpath of valdata')
parser.add_argument('--txtpath2', type=str, default=rootpath + '/test/label.txt', help='path pf val_list')
parser.add_argument('--batch_size', type=int, default=32, help='batch size for testing')
args = parser.parse_args()
test_dataset = MyDataset(args.root2, args.txtpath2, transform=None)
test_loader = DataLoader(dataset=test_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0)
model = torch.load('./modelpth/68.pth')
test_dataset = MyDataset(args.root2, args.txtpath2, transform=None)
test_loader = DataLoader(dataset=test_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0)
criterion = nn.CrossEntropyLoss()
model.eval()

model = torch.load('./modelpth/68.pth')
accuracy, _, _, C, class_accuracy = test(model, test_loader, criterion, num_classes=6)
print('Overall Accuracy:', accuracy)
for i in range(len(class_accuracy)):
    print('Class {} Accuracy: {:.2f}%'.format(i, class_accuracy[i] * 100))