pytorch 图像分割 PR曲线 代码

以下是PyTorch实现的图像分割PR曲线代码示例：

import torch
import torch.nn.functional as F
from sklearn.metrics import precision_recall_curve
import matplotlib.pyplot as plt

def plot_pr_curve(model, dataloader):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    model.eval()
    true_labels = []
    pred_scores = []
    with torch.no_grad():
        for images, labels in dataloader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            preds = F.softmax(outputs, dim=1)[:, 1] # 取出正类的概率
            true_labels += labels.cpu().numpy().tolist()
            pred_scores += preds.cpu().numpy().tolist()
    precision, recall, thresholds = precision_recall_curve(true_labels, pred_scores)
    plt.plot(recall, precision)
    plt.xlabel('Recall')
    plt.ylabel('Precision')
    plt.title('PR Curve')
    plt.show()

这个函数接收一个已经训练好的PyTorch模型和一个数据加载器，然后计算模型在数据集上的PR曲线并绘制出来。需要注意的是，这个函数假设模型的输出是一个二分类问题，因此它取出了输出中正类的概率，并将它们用于计算PR曲线。如果你的模型输出不是二分类问题，那么需要根据实际情况修改代码。