[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()就会画在一张图上

巴特西

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

模型测试

评价指标

混淆矩阵

ROC曲线图

最新文章

热门文章