import torch
import numpy as np
from sklearn.metrics import roc_auc_score
def multi_label_auc(y_true, y_pred):
"""
计算多标签AUC的函数
:param y_true: 真实标签,形状为[N, num_classes]
:param y_pred: 预测标签,形状为[N, num_classes]
:return: 多标签AUC
"""
# 将标签转换为numpy数组
y_true = y_true.cpu().numpy()
y_pred = y_pred.cpu().numpy()
# 初始化多标签AUC值
total_auc = 0.
# 计算每个标签的AUC值,并对所有标签的AUC值求平均
for i in range(y_true.shape[1]):
try:
auc = roc_auc_score(y_true[:, i], y_pred[:, i])
except ValueError:
auc = 0.5 # 如果标签中只有一个类别,则返回0.5
total_auc += auc
multi_auc = total_auc / y_true.shape[1]
return multi_auc
使用sklearn.metrics库中的roc_auc_score函数计算每个标签的AUC值,并对所有标签的AUC值求平均,从而得到多标签AUC值。在计算AUC之前,我们将PyTorch张量转换为NumPy数组。请注意,此代码假定每个标签都是二进制标签。如果您的标签具有多个类别,则需要更改AUC计算的方式。
示例:
# 随机生成的标签和预测值可能不准确
import torch
# 假设有3个标签和10个样本
num_classes = 3
N = 10
# 随机生成真实标签和预测标签
y_true = torch.randint(0, 2, size=(N, num_classes)).float()
y_pred = torch.randn(N, num_classes)
# 计算多标签AUC
multi_auc = multi_label_auc(y_true, y_pred)
print('多标签AUC:', multi_auc)
运行结果: