Move sample_pdf into PyTorch3D

Summary: Copy the sample_pdf operation from the NeRF project in to PyTorch3D, in preparation for optimizing it.

Reviewed By: gkioxari

Differential Revision: D27117930

fbshipit-source-id: 20286b007f589a4c4d53ed818c4bc5f2abd22833

tests/bm_sample_pdf.py

@@ -0,0 +1,37 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from itertools import product
+
+from fvcore.common.benchmark import benchmark
+from test_sample_pdf import TestSamplePDF
+
+
+def bm_sample_pdf() -> None:
+
+    backends = ["python_cuda", "python_cpu"]
+
+    kwargs_list = []
+    sample_counts = [64]
+    batch_sizes = [1024, 10240]
+    bin_counts = [62, 600]
+    test_cases = product(backends, sample_counts, batch_sizes, bin_counts)
+    for case in test_cases:
+        backend, n_samples, batch_size, n_bins = case
+        kwargs_list.append(
+            {
+                "backend": backend,
+                "n_samples": n_samples,
+                "batch_size": batch_size,
+                "n_bins": n_bins,
+            }
+        )
+
+    benchmark(TestSamplePDF.bm_fn, "SAMPLE_PDF", kwargs_list, warmup_iters=1)
+
+
+if __name__ == "__main__":
+    bm_sample_pdf()

tests/test_sample_pdf.py

@@ -0,0 +1,42 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from common_testing import TestCaseMixin
+from pytorch3d.renderer.implicit.sample_pdf import sample_pdf_python
+
+
+class TestSamplePDF(TestCaseMixin, unittest.TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        torch.manual_seed(1)
+
+    def test_single_bin(self):
+        bins = torch.arange(2).expand(5, 2) + 17
+        weights = torch.ones(5, 1)
+        output = sample_pdf_python(bins, weights, 100, True)
+        calc = torch.linspace(17, 18, 100).expand(5, -1)
+        self.assertClose(output, calc)
+
+    @staticmethod
+    def bm_fn(*, backend: str, n_samples, batch_size, n_bins):
+        f = sample_pdf_python
+        weights = torch.rand(size=(batch_size, n_bins))
+        bins = torch.cumsum(torch.rand(size=(batch_size, n_bins + 1)), dim=-1)
+
+        if "cuda" in backend:
+            weights = weights.cuda()
+            bins = bins.cuda()
+
+        torch.cuda.synchronize()
+
+        def output():
+            f(bins, weights, n_samples)
+            torch.cuda.synchronize()
+
+        return output