Adding support for pytorch tensors in CPU/GPU

example.py

@@ -0,0 +1,45 @@
+import numpy as np
+import torch
+
+from time import process_time
+from prdc import compute_prdc, compute_prdc_torch
+
+
+num_real_samples = num_fake_samples = 10000
+feature_dim = 1000
+nearest_k = 5
+real_features = np.random.normal(loc=0.0, scale=1.0,
+                                 size=[num_real_samples, feature_dim])
+
+fake_features = np.random.normal(loc=0.0, scale=1.0,
+                                 size=[num_fake_samples, feature_dim])
+real_features_torch = torch.Tensor(real_features)
+fake_features_torch = torch.Tensor(fake_features)
+
+start_time = process_time()
+metrics = compute_prdc(real_features=real_features,
+                       fake_features=fake_features,
+                       nearest_k=nearest_k)
+print('total time (numpy/CPU): ' + str(process_time() - start_time))
+print(metrics)
+
+start_time = process_time()
+metrics = compute_prdc_torch(real_features=real_features_torch,
+                       fake_features=fake_features_torch,
+                       nearest_k=nearest_k)
+print('total time (torch/CPU): ' + str(process_time() - start_time))
+print(metrics)
+
+if torch.cuda.is_available():
+    print('profiling in GPU')
+    real_features_torch_gpu = real_features_torch.cuda()
+    fake_features_torch_gpu = fake_features_torch.cuda()
+    start_time = process_time()
+    metrics = compute_prdc_torch(real_features=real_features_torch_gpu,
+                           fake_features=fake_features_torch_gpu,
+                           nearest_k=nearest_k)
+    torch.cuda.synchronize()
+    print('total time (torch/GPU): ' + str(process_time() - start_time))
+    print(metrics)
+else:
+    print('cuda GPU not available')

prdc/__init__.py

@@ -1 +1,7 @@
 from .prdc import compute_prdc
+
+try:
+    from .prdc_torch import compute_prdc_torch
+except ModuleNotFoundError:
+    # Error handling
+    pass

prdc/prdc_torch.py

@@ -0,0 +1,58 @@
+import torch
+from typing import Dict
+
+
+def nearest_neighbour_distances(input_features, nearest_k):
+    """
+    Args:
+        input_features: torch.Tensor([N, feature_dim], dtype=np.float32)
+        nearest_k: int
+    Returns:
+        Distances to kth nearest neighbours.
+    """
+    distances = torch.cdist(input_features, input_features)
+    radii = torch.kthvalue(distances, k=nearest_k + 1, dim=-1)[0]
+    return radii
+
+
+def compute_prdc_torch(real_features, fake_features, nearest_k):
+    """
+    Computes precision, recall, density, and coverage given two manifolds.
+
+    Args:
+        real_features: torch.Tensor([N, feature_dim], dtype=torch.float32)
+        fake_features: torch.Tensor([N, feature_dim], dtype=torch.float32)
+        nearest_k: int.
+    Returns:
+        dict of precision, recall, density, and coverage.
+    """
+
+    print('Num real: {} Num fake: {}'
+          .format(real_features.shape[0], fake_features.shape[0]))
+
+    real_nearest_neighbour_distances = nearest_neighbour_distances(
+        real_features, nearest_k)
+    fake_nearest_neighbour_distances = nearest_neighbour_distances(
+        fake_features, nearest_k)
+    distance_real_fake = torch.cdist(
+            real_features, fake_features)
+
+    precision = (
+            distance_real_fake < real_nearest_neighbour_distances[:, None]
+    ).any(dim=0).double().mean().item()
+
+    recall = (
+            distance_real_fake < fake_nearest_neighbour_distances[None, :]
+    ).any(dim=1).double().mean().item()
+
+    density = (1. / float(nearest_k)) * (
+            distance_real_fake < real_nearest_neighbour_distances[:, None]
+    ).sum(dim=0).double().mean().item()
+
+    coverage = (
+            distance_real_fake.min(dim=1)[0] <
+            real_nearest_neighbour_distances
+    ).double().mean().item()
+
+    return dict(precision=precision, recall=recall,
+                density=density, coverage=coverage)