Skip to content

Commit 8ffedde

Browse files
yifanfeng97yifanfeng97
committed
update structure vis and add RetrievalEvaluator class
1 parent 9e93f29 commit 8ffedde

9 files changed

Lines changed: 136 additions & 100 deletions

File tree

‎dhg/metrics/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
available_classification_metrics,
44
VertexClassificationEvaluator,
55
)
6-
from .retrieval import available_retrieval_metrics
6+
from .retrieval import available_retrieval_metrics, RetrievalEvaluator
77
from .recommender import available_recommender_metrics, UserItemRecommenderEvaluator
88
from .graphs import GraphVertexClassificationEvaluator
99
from .hypergraphs import HypergraphVertexClassificationEvaluator
@@ -46,4 +46,5 @@ def build_evaluator(
4646
"GraphVertexClassificationEvaluator",
4747
"HypergraphVertexClassificationEvaluator",
4848
"UserItemRecommenderEvaluator",
49+
"RetrievalEvaluator",
4950
]

‎dhg/metrics/retrieval.py‎

Lines changed: 87 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
1-
from typing import Optional, Union, Tuple
1+
from typing import Optional, Union, Tuple, List, Dict
22

33
import torch
44
import numpy as np
55

6+
from dhg.metrics.base import BaseEvaluator
7+
68

79
__all__ = [
810
"available_retrieval_metrics",
@@ -34,9 +36,7 @@ def _format_inputs(
3436
``y_pred`` (``torch.Tensor``): A 1-D tensor or 2-D tensor. Size :math:`(N_{target},)` or :math:`(N_{samples}, N_{target})`.
3537
``k`` (``int``, optional): The specified top-k value. Default to :math:`N_{target}`.
3638
"""
37-
assert (
38-
y_true.shape == y_pred.shape
39-
), "The shape of y_true and y_pred must be the same."
39+
assert y_true.shape == y_pred.shape, "The shape of y_true and y_pred must be the same."
4040
assert y_true.dim() in (1, 2), "The input y_true must be 1-D or 2-D."
4141
assert y_pred.dim() in (1, 2), "The input y_pred must be 1-D or 2-D."
4242
if y_true.dim() == 1:
@@ -50,10 +50,7 @@ def _format_inputs(
5050

5151

5252
def precision(
53-
y_true: torch.Tensor,
54-
y_pred: torch.Tensor,
55-
k: Optional[int] = None,
56-
ret_batch: bool = False,
53+
y_true: torch.Tensor, y_pred: torch.Tensor, k: Optional[int] = None, ret_batch: bool = False,
5754
) -> Union[float, list]:
5855
r"""Calculate the Precision score for the retrieval task.
5956
@@ -82,10 +79,7 @@ def precision(
8279

8380

8481
def recall(
85-
y_true: torch.Tensor,
86-
y_pred: torch.Tensor,
87-
k: Optional[int] = None,
88-
ret_batch: bool = False,
82+
y_true: torch.Tensor, y_pred: torch.Tensor, k: Optional[int] = None, ret_batch: bool = False,
8983
) -> Union[float, list]:
9084
r"""Calculate the Recall score for the retrieval task.
9185
@@ -116,10 +110,7 @@ def recall(
116110

117111

118112
def ap(
119-
y_true: torch.Tensor,
120-
y_pred: torch.Tensor,
121-
k: Optional[int] = None,
122-
method: str = "pascal_voc",
113+
y_true: torch.Tensor, y_pred: torch.Tensor, k: Optional[int] = None, method: str = "pascal_voc",
123114
) -> Union[float, list]:
124115
r"""Calculate the Average Precision (AP) for the retrieval task.
125116
@@ -137,13 +128,8 @@ def ap(
137128
>>> dm.retrieval.ap(y_true, y_pred, method="legacy")
138129
0.8333333730697632
139130
"""
140-
assert method in (
141-
"legacy",
142-
"pascal_voc",
143-
), "The method must be either legacy or pascal_voc."
144-
assert (
145-
y_true.shape == y_pred.shape
146-
), "The shape of y_true and y_pred must be the same."
131+
assert method in ("legacy", "pascal_voc",), "The method must be either legacy or pascal_voc."
132+
assert y_true.shape == y_pred.shape, "The shape of y_true and y_pred must be the same."
147133
assert y_true.dim() == 1, "The input y_true must be 1-D."
148134
assert y_pred.dim() == 1, "The input y_pred must be 1-D."
149135
y_true, y_pred = y_true.detach().float(), y_pred.detach().float()
@@ -192,10 +178,7 @@ def map(
192178
>>> dm.retrieval.map(y_true, y_pred, method="legacy")
193179
0.587037056684494
194180
"""
195-
assert method in (
196-
"legacy",
197-
"pascal_voc",
198-
), "The method must be either legacy or pascal_voc."
181+
assert method in ("legacy", "pascal_voc",), "The method must be either legacy or pascal_voc."
199182
y_true, y_pred, k = _format_inputs(y_true, y_pred, k)
200183
res_list = [ap(y_true[i, :], y_pred[i, :], k) for i in range(y_true.shape[0])]
201184
if ret_batch:
@@ -212,17 +195,12 @@ def _dcg(matrix: torch.Tensor) -> torch.Tensor:
212195
"""
213196
assert matrix.dim() == 2, "The input must be a 2-D tensor."
214197
n, k = matrix.shape
215-
denom = (
216-
torch.log2(torch.arange(k, device=matrix.device) + 2.0).view(1, -1).repeat(n, 1)
217-
)
198+
denom = torch.log2(torch.arange(k, device=matrix.device) + 2.0).view(1, -1).repeat(n, 1)
218199
return (matrix / denom).sum(dim=-1)
219200

220201

221202
def ndcg(
222-
y_true: torch.Tensor,
223-
y_pred: torch.Tensor,
224-
k: Optional[int] = None,
225-
ret_batch: bool = False,
203+
y_true: torch.Tensor, y_pred: torch.Tensor, k: Optional[int] = None, ret_batch: bool = False,
226204
) -> Union[float, list]:
227205
r"""Calculate the Normalized Discounted Cumulative Gain (NDCG) for the retrieval task.
228206
@@ -290,10 +268,7 @@ def rr(y_true: torch.Tensor, y_pred: torch.Tensor, k: Optional[int] = None) -> f
290268

291269

292270
def mrr(
293-
y_true: torch.Tensor,
294-
y_pred: torch.Tensor,
295-
k: Optional[int] = None,
296-
ret_batch: bool = False,
271+
y_true: torch.Tensor, y_pred: torch.Tensor, k: Optional[int] = None, ret_batch: bool = False,
297272
) -> Union[float, list]:
298273
r"""Calculate the mean Reciprocal Rank (MRR) for the retrieval task.
299274
@@ -348,13 +323,8 @@ def _pr_curve(
348323
>>> recall_coor
349324
[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
350325
"""
351-
assert method in (
352-
"legacy",
353-
"pascal_voc",
354-
), "The method must be either legacy or pascal_voc."
355-
assert (
356-
y_true.shape == y_pred.shape
357-
), "The shape of y_true and y_pred must be the same."
326+
assert method in ("legacy", "pascal_voc",), "The method must be either legacy or pascal_voc."
327+
assert y_true.shape == y_pred.shape, "The shape of y_true and y_pred must be the same."
358328
assert y_true.dim() == 1, "The input y_true must be 1-D."
359329
assert y_pred.dim() == 1, "The input y_pred must be 1-D."
360330
y_true, y_pred = y_true.detach().float(), y_pred.detach().float()
@@ -405,16 +375,11 @@ def pr_curve(
405375
>>> recall_coor
406376
[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
407377
"""
408-
assert method in (
409-
"legacy",
410-
"pascal_voc",
411-
), "The method must be either legacy or pascal_voc."
378+
assert method in ("legacy", "pascal_voc",), "The method must be either legacy or pascal_voc."
412379
y_true, y_pred, k = _format_inputs(y_true, y_pred, k)
413380
precision_coor_list, recall_coor_list = [], []
414381
for i in range(y_true.shape[0]):
415-
precision_coor, recall_coor = _pr_curve(
416-
y_true[i, :], y_pred[i, :], k, method, n_points
417-
)
382+
precision_coor, recall_coor = _pr_curve(y_true[i, :], y_pred[i, :], k, method, n_points)
418383
precision_coor_list.append(precision_coor)
419384
recall_coor_list.append(recall_coor)
420385
if ret_batch:
@@ -423,3 +388,73 @@ def pr_curve(
423388
precision_coor = np.mean(precision_coor_list, axis=0)
424389
recall_coor = np.mean(recall_coor_list, axis=0)
425390
return precision_coor.tolist(), recall_coor.tolist()
391+
392+
393+
class RetrievalEvaluator(BaseEvaluator):
394+
r"""Return the metric evaluator for retrieval task. The supported metrics includes: ``precision``, ``recall``, ``map``, ``ndcg``, ``mrr``, ``pr_curve``.
395+
396+
Args:
397+
``metric_configs`` (``List[Union[str, Dict[str, dict]]]``): The metric configurations. The key is the metric name and the value is the metric parameters.
398+
``validate_index`` (``int``): The specified metric index used for validation. Defaults to ``0``.
399+
400+
Examples:
401+
>>> import torch
402+
>>> import dhg.metrics as dm
403+
>>> evaluator = dm.RetrievalEvaluator(
404+
[
405+
"precision",
406+
"recall",
407+
"ndcg",
408+
],
409+
0
410+
)
411+
>>> y_true = torch.tensor([0, 1, 0, 0, 1, 1])
412+
>>> y_pred = torch.tensor([0.8, 0.9, 0.6, 0.7, 0.4, 0.5])
413+
>>> evaluator.validate_add_batch(y_true, y_pred)
414+
>>> y_true = torch.tensor([0, 1, 0, 1, 0, 1])
415+
>>> y_pred = torch.tensor([0.8, 0.9, 0.9, 0.4, 0.4, 0.5])
416+
>>> evaluator.validate_add_batch(y_true, y_pred)
417+
>>> evaluator.validate_epoch_res()
418+
0.5
419+
>>> y_true = torch.tensor([0, 1, 1, 1, 0, 1])
420+
>>> y_pred = torch.tensor([0.8, 0.9, 0.6, 0.7, 0.4, 0.5])
421+
>>> evaluator.test_add_batch(y_true, y_pred)
422+
>>> y_true = torch.tensor([1, 1, 0, 0, 1, 0])
423+
>>> y_pred = torch.tensor([0.8, 0.9, 0.9, 0.4, 0.4, 0.5])
424+
>>> evaluator.test_add_batch(y_true, y_pred)
425+
>>> evaluator.test_epoch_res()
426+
{'precision': 0.5833333432674408, 'recall': 1.0, 'ndcg': 0.8878978490829468}
427+
"""
428+
429+
def __init__(
430+
self, metric_configs: List[Union[str, Dict[str, dict]]], validate_index: int = 0,
431+
):
432+
super().__init__("retrieval", metric_configs, validate_index)
433+
434+
def validate_add_batch(self, batch_y_true: torch.Tensor, batch_y_pred: torch.Tensor):
435+
r"""Add batch data for validation.
436+
437+
Args:
438+
``batch_y_true`` (``torch.Tensor``): The ground truth data. Size :math:`(N_{batch}, -)`.
439+
``batch_y_pred`` (``torch.Tensor``): The predicted data. Size :math:`(N_{batch}, -)`.
440+
"""
441+
return super().validate_add_batch(batch_y_true, batch_y_pred)
442+
443+
def validate_epoch_res(self):
444+
r"""For all added batch data, return the result of the evaluation on the specified ``validate_index``-th metric.
445+
"""
446+
return super().validate_epoch_res()
447+
448+
def test_add_batch(self, batch_y_true: torch.Tensor, batch_y_pred: torch.Tensor):
449+
r"""Add batch data for testing.
450+
451+
Args:
452+
``batch_y_true`` (``torch.Tensor``): The ground truth data. Size :math:`(N_{batch}, -)`.
453+
``batch_y_pred`` (``torch.Tensor``): The predicted data. Size :math:`(N_{batch}, -)`.
454+
"""
455+
return super().test_add_batch(batch_y_true, batch_y_pred)
456+
457+
def test_epoch_res(self):
458+
r"""For all added batch data, return results of the evaluation on all the metrics in ``metric_configs``.
459+
"""
460+
return super().test_epoch_res()

‎dhg/random/hypergraphs/hypergraph.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ def uniform_hypergraph_Gnm(k: int, num_v: int, num_e: int):
6262
return Hypergraph(num_v, list(edges))
6363

6464

65-
def hypergraph_Gnm(num_v: int, num_e: int, method:str="uniform", prob_k_list: Optional[List[float]] = None):
65+
def hypergraph_Gnm(num_v: int, num_e: int, method:str="low_order_first", prob_k_list: Optional[List[float]] = None):
6666
r"""Return a random hypergraph with ``num_v`` vertices and ``num_e`` hyperedges. The ``method`` argument determines the distribution of the hyperedge degree.
6767
The ``method`` can be one of ``"uniform"``, ``"low_order_first"``, ``"high_order_first"``.
6868
@@ -90,11 +90,11 @@ def hypergraph_Gnm(num_v: int, num_e: int, method:str="uniform", prob_k_list: Op
9090
if method == "uniform":
9191
prob_k_list = [C(num_v, k) / (2 ** num_v - 1) for k in deg_e_list]
9292
elif method == "low_order_first":
93-
prob_k_list = [1.8 ** (-k) for k in deg_e_list]
93+
prob_k_list = [3 ** (-k) for k in range(len(deg_e_list))]
9494
sum_of_prob_k_list = sum(prob_k_list)
9595
prob_k_list = [prob_k / sum_of_prob_k_list for prob_k in prob_k_list]
9696
elif method == "high_order_first":
97-
prob_k_list = [1.8 ** (-k) for k in deg_e_list].reverse()
97+
prob_k_list = [3 ** (-k) for k in range(len(deg_e_list))].reverse()
9898
sum_of_prob_k_list = sum(prob_k_list)
9999
prob_k_list = [prob_k / sum_of_prob_k_list for prob_k in prob_k_list]
100100
else:

‎docs/source/_static/img/vis_bigraph.png‎

-45.7 KB
Loading

‎docs/source/_static/img/vis_digraph.png‎

-23.9 KB
Loading

‎docs/source/_static/img/vis_graph.png‎

-17.4 KB
Loading

‎docs/source/_static/img/vis_hypergraph.png‎

-12.3 KB
Loading

‎docs/source/api/metrics.rst‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,5 +90,9 @@ On User-Item Bipartite Graphs
9090
:members:
9191
:show-inheritance:
9292

93-
.. Retrieval
94-
.. +++++++++++++++++++
93+
Retrieval
94+
+++++++++++++++++++
95+
96+
.. autoclass:: dhg.metrics.RetrievalEvaluator
97+
:members:
98+
:show-inheritance:

0 commit comments

Comments
 (0)