cuda

Author: Janos95

benchmarks/benchmark_chamfer.py

@@ -6,6 +6,10 @@
 
 import chamfer
 
+try:
+    from pytorch3d.loss import chamfer_distance as pytorch3d_chamfer
+except ImportError:  # pragma: no cover - optional dependency
+    pytorch3d_chamfer = None
 
 def chunked_brute_force(query: torch.Tensor, reference: torch.Tensor, chunk: int = 1024) -> Tuple[torch.Tensor, torch.Tensor]:
     """Memory-friendly brute-force NN by processing reference points in chunks."""
@@ -65,6 +69,11 @@ def mps_sync() -> None:
         torch.mps.synchronize()
 
 
+def cuda_sync() -> None:
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(description="Benchmark chamfer nearest neighbour implementation.")
     parser.add_argument("--n", type=int, default=5_000, help="Number of points per set")
@@ -77,84 +86,183 @@ def main() -> None:
     a_cpu = torch.rand(args.n, args.dims)
     b_cpu = torch.rand(args.n, args.dims)
     mps_available = torch.backends.mps.is_available()
+    cuda_available = torch.cuda.is_available()
 
-    # Warmups
-    chamfer.closest_points(a_cpu[:256], b_cpu[:256], use_mps=False)
-    chunked_brute_force(a_cpu[:512], b_cpu[:512], chunk=args.chunk)
+    # Warmups to trigger compilation/allocation outside timing loops.
+    chunked_chamfer_loss(a_cpu[:512], b_cpu[:512], chunk=args.chunk)
+    chamfer.chamfer_distance(a_cpu[:256], b_cpu[:256], use_mps=False)
+    if pytorch3d_chamfer is not None:
+        pytorch3d_chamfer(a_cpu[:256].unsqueeze(0), b_cpu[:256].unsqueeze(0))
 
     a_mps = b_mps = None
     if mps_available:
         a_mps = a_cpu.to("mps")
         b_mps = b_cpu.to("mps")
-        chamfer.closest_points(a_mps[:256], b_mps[:256], use_mps=True)
-
-    brute_fwd_time = time_call(lambda: chunked_brute_force(a_cpu, b_cpu, chunk=args.chunk), repeat=args.repeat)
-    cpu_kd_time = time_call(lambda: chamfer.closest_points(a_cpu, b_cpu, use_mps=False), repeat=args.repeat)
-
-    brute_grad_time = time_call(
-        lambda: chunked_chamfer_loss(
-            a_cpu.clone().requires_grad_(True),
-            b_cpu.clone().requires_grad_(True),
-            chunk=args.chunk,
-        ).backward(),
-        repeat=args.repeat,
-    )
+        chamfer.chamfer_distance(a_mps[:256], b_mps[:256], use_mps=True)
+
+    a_cuda = b_cuda = None
+    if cuda_available:
+        a_cuda = a_cpu.to("cuda")
+        b_cuda = b_cpu.to("cuda")
+        chamfer.chamfer_distance(a_cuda[:256], b_cuda[:256])
+        if pytorch3d_chamfer is not None:
+            pytorch3d_chamfer(a_cuda[:256].unsqueeze(0), b_cuda[:256].unsqueeze(0))
+        cuda_sync()
+
+    def brute_forward() -> None:
+        chunked_chamfer_loss(a_cpu, b_cpu, chunk=args.chunk)
+
+    def brute_backward() -> None:
+        a = a_cpu.clone().requires_grad_(True)
+        b = b_cpu.clone().requires_grad_(True)
+        loss = chunked_chamfer_loss(a, b, chunk=args.chunk)
+        loss.backward()
+
+    brute_forward_time = time_call(brute_forward, repeat=args.repeat)
+    brute_backward_time = time_call(brute_backward, repeat=args.repeat)
 
-    def cpu_grad() -> None:
+    def kd_cpu_forward() -> None:
+        chamfer.chamfer_distance(a_cpu, b_cpu, use_mps=False)
+
+    def kd_cpu_backward() -> None:
         a = a_cpu.clone().requires_grad_(True)
         b = b_cpu.clone().requires_grad_(True)
         loss = chamfer.chamfer_distance(a, b, use_mps=False)
         loss.backward()
 
-    cpu_grad_time = time_call(cpu_grad, repeat=args.repeat)
+    cpu_forward_time = time_call(kd_cpu_forward, repeat=args.repeat)
+    cpu_backward_time = time_call(kd_cpu_backward, repeat=args.repeat)
+
+    kd_cuda_forward_time = None
+    kd_cuda_backward_time = None
+    pytorch3d_cuda_forward_time = None
+    pytorch3d_cuda_backward_time = None
+    if cuda_available and a_cuda is not None and b_cuda is not None:
+        def kd_cuda_forward() -> None:
+            chamfer.chamfer_distance(a_cuda, b_cuda)
+
+        kd_cuda_forward_time = time_call(kd_cuda_forward, sync=cuda_sync, repeat=args.repeat)
+
+        def kd_cuda_backward() -> None:
+            a = a_cuda.clone().requires_grad_(True)
+            b = b_cuda.clone().requires_grad_(True)
+            loss = chamfer.chamfer_distance(a, b)
+            loss.backward()
+
+        kd_cuda_backward_time = time_call(kd_cuda_backward, sync=cuda_sync, repeat=args.repeat)
+
+        if pytorch3d_chamfer is not None:
+            def pyt3d_cuda_forward() -> None:
+                loss, _ = pytorch3d_chamfer(a_cuda.unsqueeze(0), b_cuda.unsqueeze(0))
+                return loss
+
+            pyt3d_cuda_forward_time = time_call(pyt3d_cuda_forward, sync=cuda_sync, repeat=args.repeat)
+
+            def pyt3d_cuda_backward() -> None:
+                a = a_cuda.unsqueeze(0).clone().requires_grad_(True)
+                b = b_cuda.unsqueeze(0).clone().requires_grad_(True)
+                loss, _ = pytorch3d_chamfer(a, b)
+                loss.backward()
+
+            pyt3d_cuda_backward_time = time_call(pyt3d_cuda_backward, sync=cuda_sync, repeat=args.repeat)
 
-    kd_mps_time = None
-    mps_grad_time = None
+    kd_mps_forward_time = None
+    kd_mps_backward_time = None
     if mps_available and a_mps is not None and b_mps is not None:
-        kd_mps_time = time_call(
-            lambda: chamfer.closest_points(a_mps, b_mps, use_mps=True),
-            sync=mps_sync,
-            repeat=args.repeat,
-        )
+        def kd_mps_forward() -> None:
+            chamfer.chamfer_distance(a_mps, b_mps, use_mps=True)
+
+        kd_mps_forward_time = time_call(kd_mps_forward, sync=mps_sync, repeat=args.repeat)
 
-        def mps_grad() -> None:
+        def kd_mps_backward() -> None:
             a = a_mps.clone().requires_grad_(True)
             b = b_mps.clone().requires_grad_(True)
             loss = chamfer.chamfer_distance(a, b, use_mps=True)
             loss.backward()
 
-        mps_grad_time = time_call(mps_grad, sync=mps_sync, repeat=args.repeat)
+        kd_mps_backward_time = time_call(kd_mps_backward, sync=mps_sync, repeat=args.repeat)
+
+    pyt3d_cpu_forward_time = None
+    pyt3d_cpu_backward_time = None
+    if pytorch3d_chamfer is not None:
+        def pyt3d_cpu_forward() -> None:
+            loss, _ = pytorch3d_chamfer(a_cpu.unsqueeze(0), b_cpu.unsqueeze(0))
+            return loss
+
+        pyt3d_cpu_forward_time = time_call(pyt3d_cpu_forward, repeat=args.repeat)
+
+        def pyt3d_cpu_backward() -> None:
+            a = a_cpu.unsqueeze(0).clone().requires_grad_(True)
+            b = b_cpu.unsqueeze(0).clone().requires_grad_(True)
+            loss, _ = pytorch3d_chamfer(a, b)
+            loss.backward()
+
+        pyt3d_cpu_backward_time = time_call(pyt3d_cpu_backward, repeat=args.repeat)
 
     # Prepare table rows
     rows = []
 
     rows.append(
         (
             "Brute force",
-            f"{brute_fwd_time:.3f} s",
-            f"{brute_grad_time:.3f} s",
+            f"{brute_forward_time:.3f} s",
+            f"{brute_backward_time:.3f} s",
         )
     )
 
     rows.append(
         (
             "KD-tree CPU",
-            f"{cpu_kd_time:.3f} s ({brute_fwd_time / cpu_kd_time:.2f}x)",
-            f"{cpu_grad_time:.3f} s ({brute_grad_time / cpu_grad_time:.2f}x)",
+            f"{cpu_forward_time:.3f} s ({brute_forward_time / cpu_forward_time:.2f}x)",
+            f"{cpu_backward_time:.3f} s ({brute_backward_time / cpu_backward_time:.2f}x)",
         )
     )
 
-    if kd_mps_time is not None and mps_grad_time is not None:
+    if kd_cuda_forward_time is not None and kd_cuda_backward_time is not None:
+        rows.append(
+            (
+                "KD-tree CUDA",
+                f"{kd_cuda_forward_time:.3f} s ({brute_forward_time / kd_cuda_forward_time:.2f}x)",
+                f"{kd_cuda_backward_time:.3f} s ({brute_backward_time / kd_cuda_backward_time:.2f}x)",
+            )
+        )
+    else:
+        rows.append(("KD-tree CUDA", "n/a", "n/a"))
+
+    if kd_mps_forward_time is not None and kd_mps_backward_time is not None:
         rows.append(
             (
                 "KD-tree MPS",
-                f"{kd_mps_time:.3f} s ({brute_fwd_time / kd_mps_time:.2f}x)",
-                f"{mps_grad_time:.3f} s ({brute_grad_time / mps_grad_time:.2f}x)",
+                f"{kd_mps_forward_time:.3f} s ({brute_forward_time / kd_mps_forward_time:.2f}x)",
+                f"{kd_mps_backward_time:.3f} s ({brute_backward_time / kd_mps_backward_time:.2f}x)",
             )
         )
     else:
         rows.append(("KD-tree MPS", "n/a", "n/a"))
 
+    if pytorch3d_chamfer is not None and pyt3d_cpu_forward_time is not None and pyt3d_cpu_backward_time is not None:
+        rows.append(
+            (
+                "PyTorch3D CPU",
+                f"{pyt3d_cpu_forward_time:.3f} s ({brute_forward_time / pyt3d_cpu_forward_time:.2f}x)",
+                f"{pyt3d_cpu_backward_time:.3f} s ({brute_backward_time / pyt3d_cpu_backward_time:.2f}x)",
+            )
+        )
+
+        if (
+            pyt3d_cuda_forward_time is not None
+            and pyt3d_cuda_backward_time is not None
+        ):
+            rows.append(
+                (
+                    "PyTorch3D CUDA",
+                    f"{pyt3d_cuda_forward_time:.3f} s ({brute_forward_time / pyt3d_cuda_forward_time:.2f}x)",
+                    f"{pyt3d_cuda_backward_time:.3f} s ({brute_backward_time / pyt3d_cuda_backward_time:.2f}x)",
+                )
+            )
+        else:
+            rows.append(("PyTorch3D CUDA", "n/a", "n/a"))
+
     header = ("Method", "Forward", "Backward")
     widths = [max(len(col), max(len(row[i]) for row in rows)) for i, col in enumerate(header)]