From 4878a0abe50e14aa04c220d601208935cf04069e Mon Sep 17 00:00:00 2001 From: Frozenmad Date: Mon, 31 May 2021 13:19:39 +0800 Subject: [PATCH] adjust metrics input --- autogl/module/train/evaluation.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/autogl/module/train/evaluation.py b/autogl/module/train/evaluation.py index af3a79a..d36f708 100644 --- a/autogl/module/train/evaluation.py +++ b/autogl/module/train/evaluation.py @@ -120,6 +120,7 @@ class Auc(Evaluation): if len(predict.shape) == 1: pos_predict = predict else: + assert predict.shape[1] == 2, "Cannot use auc on given data with %d classes!" % (predict.shape[1]) pos_predict = predict[:, 1] return roc_auc_score(label, pos_predict) @@ -142,7 +143,11 @@ class Acc(Evaluation): """ Should return: the evaluation result (float) """ - return accuracy_score(label, np.argmax(predict, axis=1)) + if len(predict.shape) == 2: + predict = np.argmax(predict, axis=1) + else: + predict = [1 if p > 0.5 else 0 for p in predict] + return accuracy_score(label, predict) @register_evaluate("mrr") @@ -163,5 +168,9 @@ class Mrr(Evaluation): """ Should return: the evaluation result (float) """ - pos_predict = predict[:, 1] + if len(predict.shape) == 2: + assert predict.shape[1] == 2, "Cannot use mrr on given data with %d classes!" % (predict.shape[1]) + pos_predict = predict[:, 1] + else: + pos_predict = predict return label_ranking_average_precision_score(label, pos_predict)