网站首页 > 厂商资讯 > deepflow >

如何在PyTorch中可视化模型PR曲线？

在深度学习领域，性能评估是模型开发过程中的关键环节。其中，ROC曲线（Receiver Operating Characteristic Curve）和PR曲线（Precision-Recall Curve）是两种常用的性能评估方法。本文将详细介绍如何在PyTorch中可视化模型的PR曲线，帮助您更好地理解模型性能。

一、PR曲线简介

PR曲线是评估二分类模型性能的一种图形化方法，它展示了模型在不同阈值下的精确率（Precision）和召回率（Recall）之间的关系。与ROC曲线相比，PR曲线更适合于不平衡数据集，因为它更关注精确率。

二、PyTorch中可视化PR曲线的步骤

  1. 准备数据集

    首先，您需要准备一个二分类数据集。这里以MNIST手写数字数据集为例。

    import torchvision.datasets as datasets
    
import torchvision.transforms as transforms
    

    
train_dataset = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
    
test_dataset = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor(), download=True)

  2. 定义模型

    在PyTorch中，您可以使用torch.nn模块定义模型。以下是一个简单的卷积神经网络模型：

    import torch.nn as nn
    
import torch.nn.functional as F
    

    
class CNN(nn.Module):
    
    def __init__(self):
    
        super(CNN, self).__init__()
    
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
    
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
    
        self.fc1 = nn.Linear(320, 50)
    
        self.fc2 = nn.Linear(50, 10)
    

    
    def forward(self, x):
    
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
    
        x = F.relu(F.max_pool2d(self.conv2(x), 2))
    
        x = x.view(-1, 320)
    
        x = F.relu(self.fc1(x))
    
        x = self.fc2(x)
    
        return F.log_softmax(x, dim=1)

  3. 训练模型

    使用PyTorch的torch.optim模块定义优化器，并使用torch.utils.data.DataLoader进行数据加载。

    import torch.optim as optim
    

    
model = CNN()
    
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    
criterion = nn.CrossEntropyLoss()
    

    
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
    
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)
    

    
for epoch in range(10):  # 训练10个epoch
    
    for data, target in train_loader:
    
        optimizer.zero_grad()
    
        output = model(data)
    
        loss = criterion(output, target)
    
        loss.backward()
    
        optimizer.step()

  4. 计算PR曲线

    在PyTorch中，可以使用sklearn.metrics模块计算PR曲线。

    from sklearn.metrics import precision_recall_curve
    

    
model.eval()
    
all_targets = []
    
all_scores = []
    

    
with torch.no_grad():
    
    for data, target in test_loader:
    
        output = model(data)
    
        _, predicted = torch.max(output, 1)
    
        all_targets.extend(target.numpy())
    
        all_scores.extend(output.data.cpu().numpy().flatten())
    

    
precision, recall, thresholds = precision_recall_curve(all_targets, all_scores)

  5. 可视化PR曲线

    使用matplotlib模块绘制PR曲线。

    import matplotlib.pyplot as plt
    

    
plt.figure(figsize=(8, 6))
    
plt.plot(recall, precision, label='PR curve')
    
plt.xlabel('Recall')
    
plt.ylabel('Precision')
    
plt.title('Precision-Recall Curve')
    
plt.legend()
    
plt.show()

三、案例分析

以下是一个使用PyTorch和sklearn.metrics模块计算PR曲线的完整示例：

import torch

import torchvision

import torchvision.transforms as transforms

from sklearn.metrics import precision_recall_curve

import matplotlib.pyplot as plt



# 准备数据集

train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)

test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor(), download=True)



# 定义模型

class CNN(nn.Module):

    def __init__(self):

        super(CNN, self).__init__()

        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)

        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)

        self.fc1 = nn.Linear(320, 50)

        self.fc2 = nn.Linear(50, 10)



    def forward(self, x):

        x = F.relu(F.max_pool2d(self.conv1(x), 2))

        x = F.relu(F.max_pool2d(self.conv2(x), 2))

        x = x.view(-1, 320)

        x = F.relu(self.fc1(x))

        x = self.fc2(x)

        return F.log_softmax(x, dim=1)



model = CNN()

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

criterion = nn.CrossEntropyLoss()



train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)

test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)



# 训练模型

for epoch in range(10):

    for data, target in train_loader:

        optimizer.zero_grad()

        output = model(data)

        loss = criterion(output, target)

        loss.backward()

        optimizer.step()



# 计算PR曲线

model.eval()

all_targets = []

all_scores = []



with torch.no_grad():

    for data, target in test_loader:

        output = model(data)

        _, predicted = torch.max(output, 1)

        all_targets.extend(target.numpy())

        all_scores.extend(output.data.cpu().numpy().flatten())



precision, recall, thresholds = precision_recall_curve(all_targets, all_scores)



# 可视化PR曲线

plt.figure(figsize=(8, 6))

plt.plot(recall, precision, label='PR curve')

plt.xlabel('Recall')

plt.ylabel('Precision')

plt.title('Precision-Recall Curve')

plt.legend()

plt.show()

通过以上步骤，您可以在PyTorch中可视化模型的PR曲线，从而更好地评估模型性能。在实际应用中，您可以根据自己的需求调整模型结构和参数，以达到最佳性能。

猜你喜欢：全链路监控