Remove internal metrics in favor of torchmetrics

examples/dgcnn_segmentation.py

@@ -2,12 +2,12 @@
 
 import torch
 import torch.nn.functional as F
+from torchmetrics.functional import jaccard_index
 
 import torch_geometric.transforms as T
 from torch_geometric.datasets import ShapeNet
 from torch_geometric.loader import DataLoader
 from torch_geometric.nn import MLP, DynamicEdgeConv
-from torch_geometric.utils import intersection_and_union as i_and_u
 
 category = 'Airplane'  # Pass in `None` to train on all categories.
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'ShapeNet')
@@ -86,11 +86,8 @@ def test(loader):
     for data in loader:
         data = data.to(device)
         pred = model(data).argmax(dim=1)
-
-        i, u = i_and_u(pred, data.y, loader.dataset.num_classes, data.batch)
-        iou = i.cpu().to(torch.float) / u.cpu().to(torch.float)
-        iou[torch.isnan(iou)] = 1
-
+        iou = jaccard_index(pred, data.y,
+                            num_classes=loader.dataset.num_classes)
         # Find and filter the relevant classes for each category.
         for iou, category in zip(iou.unbind(), data.category.unbind()):
             ious[category.item()].append(iou[y_mask[category]])

examples/point_transformer_segmentation.py

@@ -11,12 +11,12 @@
 from torch.nn import ReLU
 from torch.nn import Sequential as Seq
 from torch_cluster import knn_graph
+from torchmetrics.functional import jaccard_index
 
 import torch_geometric.transforms as T
 from torch_geometric.datasets import ShapeNet
 from torch_geometric.loader import DataLoader
 from torch_geometric.nn.unpool import knn_interpolate
-from torch_geometric.utils import intersection_and_union as i_and_u
 
 category = 'Airplane'  # Pass in `None` to train on all categories.
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'ShapeNet')
@@ -203,11 +203,8 @@ def test(loader):
     for data in loader:
         data = data.to(device)
         pred = model(data.x, data.pos, data.batch).argmax(dim=1)
-
-        i, u = i_and_u(pred, data.y, loader.dataset.num_classes, data.batch)
-        iou = i.cpu().to(torch.float) / u.cpu().to(torch.float)
-        iou[torch.isnan(iou)] = 1
-
+        iou = jaccard_index(pred, data.y,
+                            num_classes=loader.dataset.num_classes)
         # Find and filter the relevant classes for each category.
         for iou, category in zip(iou.unbind(), data.category.unbind()):
             ious[category.item()].append(iou[y_mask[category]])

examples/pointnet2_segmentation.py

@@ -3,12 +3,12 @@
 import torch
 import torch.nn.functional as F
 from pointnet2_classification import GlobalSAModule, SAModule
+from torchmetrics.functional import jaccard_index
 
 import torch_geometric.transforms as T
 from torch_geometric.datasets import ShapeNet
 from torch_geometric.loader import DataLoader
 from torch_geometric.nn import MLP, knn_interpolate
-from torch_geometric.utils import intersection_and_union as i_and_u
 
 category = 'Airplane'  # Pass in `None` to train on all categories.
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'ShapeNet')
@@ -112,11 +112,8 @@ def test(loader):
     for data in loader:
         data = data.to(device)
         pred = model(data).argmax(dim=1)
-
-        i, u = i_and_u(pred, data.y, loader.dataset.num_classes, data.batch)
-        iou = i.cpu().to(torch.float) / u.cpu().to(torch.float)
-        iou[torch.isnan(iou)] = 1
-
+        iou = jaccard_index(pred, data.y,
+                            num_classes=loader.dataset.num_classes)
         # Find and filter the relevant classes for each category.
         for iou, category in zip(iou.unbind(), data.category.unbind()):
             ious[category.item()].append(iou[y_mask[category]])

setup.py

@@ -28,6 +28,7 @@
     'scikit-image',
     'pytorch-memlab',
     'class-resolver>=0.3.2',
+    'torchmetrics>=0.7',
 ]
 
 test_requires = [

torch_geometric/utils/__init__.py

@@ -31,9 +31,6 @@
                                 structured_negative_sampling,
                                 structured_negative_sampling_feasible)
 from .train_test_split_edges import train_test_split_edges
-from .metric import (accuracy, true_positive, true_negative, false_positive,
-                     false_negative, precision, recall, f1_score,
-                     intersection_and_union, mean_iou)
 
 __all__ = [
     'degree',
@@ -80,16 +77,6 @@
     'structured_negative_sampling',
     'structured_negative_sampling_feasible',
     'train_test_split_edges',
-    'accuracy',
-    'true_positive',
-    'true_negative',
-    'false_positive',
-    'false_negative',
-    'precision',
-    'recall',
-    'f1_score',
-    'intersection_and_union',
-    'mean_iou',
 ]
 
 classes = __all__