mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 03:42:50 +08:00
7d7d00f288
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
43 lines
1.2 KiB
Python
43 lines
1.2 KiB
Python
# 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
|