[AI]-模型测试和评价指标

模型测试

import cv2
from torchvision import transforms, datasets, models
from torch.utils.data import  DataLoader
import torch
import numpy as np
import os
from sklearn import metrics
import matplotlib.pyplot as plt

device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")
print(device)
num_class = 3
model_path = 模型路径
model = 模型(num_class).to(device)
model.load_state_dict(torch.load(model_path))
model.eval()   # Set model to evaluate mode

test_dataset = 数据集读取(train=False)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=2)

def turn(l):
    l = l.data.cpu().numpy()
    l = l.squeeze()
    l = np.swapaxes(l, 0, 2)
    l = np.swapaxes(l, 0, 1)
    return l

for inputs, labels in test_loader:
    model.to(device)
    inputs = inputs.to(device)
    labels = labels.to(device)

    pred = model(inputs)
    # pred = torch.relu(pred)
    pred = turn(pred)
    gt = turn(labels)

评价指标

混淆矩阵

以分割为例，经过.flatten()处理。

def acc(pred, gt):
    tp = 0
    tn = 0
    fp = 0
    fn = 0
    num = len(pred)
    for i in range(num):
        if pred[i] > 0 and gt[i] == 1:
            tp += 1
        if pred[i] > 0 and gt[i] == 0:
            fp += 1
        if pred[i] == 0 and gt[i] == 1:
            fn += 1
        if pred[i] == 0 and gt[i] == 0:
            tn += 1
    acc = (tp + tn) / num
    iou = tp / (tp + fp + fn)
    rec = tp / (tp + fn)
    pre = tp / (tp + fp)
    f1 = 2 * pre * rec / (pre + rec)
    print("mAcc is :{},  mIou is :{}, recall is :{}, precision is :{}, f1 is :{}".format(acc, iou, rec, pre, f1))

ROC曲线图

def draw_roc(pred, gt, name):
    tpr, fpr, thresholds = metrics.roc_curve(gt, pred, pos_label=0)
    plt.figure
    plt.plot(fpr, tpr, label = name)
    plt.xlabel('FPR')
    plt.ylabel('TPR')
    plt.legend(loc = 'lower right')
    plt.title(name)
    plt.savefig('路径/{}.png'.format(name))
    # plt.close()  如果有多个类别，不close()就会画在一张图上