From b99c19d4a85cdc604532edc16b871b642c9bde16 Mon Sep 17 00:00:00 2001 From: Ruilong Li Date: Tue, 12 Jul 2022 06:25:55 +0000 Subject: [PATCH 1/2] speedup PDFSampler impl. --- pyrad/graphs/modules/ray_sampler.py | 62 +++++++++++++++++------------ 1 file changed, 37 insertions(+), 25 deletions(-) diff --git a/pyrad/graphs/modules/ray_sampler.py b/pyrad/graphs/modules/ray_sampler.py index 4ccccabd2..db8ccb86b 100644 --- a/pyrad/graphs/modules/ray_sampler.py +++ b/pyrad/graphs/modules/ray_sampler.py @@ -23,8 +23,8 @@ import torch from torch import nn from torchtyping import TensorType -from pyrad.fields.occupancy_fields.occupancy_grid import OccupancyGrid from pyrad.cameras.rays import RayBundle, RaySamples +from pyrad.fields.occupancy_fields.occupancy_grid import OccupancyGrid class Sampler(nn.Module): @@ -305,34 +305,34 @@ class PDFSampler(Sampler): u = torch.linspace(0.0, 1.0, steps=num_bins, device=cdf.device) u = u.expand(size=(*cdf.shape[:-1], num_bins)) - mask = u[..., None, :] >= cdf[..., :, None] # [num_samples, num_orig_bins, num_bins] + # mask = u[..., None, :] >= cdf[..., :, None] # [num_samples, num_orig_bins, num_bins] - # Uses same interval trick as mip-NeRF - def find_interval( - mask: TensorType[..., "num_orig_bins", "num_bins"], x: TensorType[..., "num_orig_bins"] - ) -> Tuple[TensorType[..., "num_bins"], TensorType[..., "num_bins"]]: - """Find intervals based on cdf mask. + # # Uses same interval trick as mip-NeRF + # def find_interval( + # mask: TensorType[..., "num_orig_bins", "num_bins"], x: TensorType[..., "num_orig_bins"] + # ) -> Tuple[TensorType[..., "num_bins"], TensorType[..., "num_bins"]]: + # """Find intervals based on cdf mask. - Mask x x_start x_end - T T T T F [x0 x1 x2 x3 x4] [x0 x2 x3 x3] [x1 x3 x4 x4] - T T T T F - T T T F F - T F F F F + # Mask x x_start x_end + # T T T T F [x0 x1 x2 x3 x4] [x0 x2 x3 x3] [x1 x3 x4 x4] + # T T T T F + # T T T F F + # T F F F F - Where the number of rows correspond to the target number of bins. The number of columns - correspond to the input number of bins + # Where the number of rows correspond to the target number of bins. The number of columns + # correspond to the input number of bins - Args: - mask (TensorType[..., "num_orig_bins", "num_bins"]): PDF represented as a boolean mask. - x (TensorType[..., "num_original_bins"]): Probe to calculate intervals for. + # Args: + # mask (TensorType[..., "num_orig_bins", "num_bins"]): PDF represented as a boolean mask. + # x (TensorType[..., "num_original_bins"]): Probe to calculate intervals for. - Returns: - Tuple[TensorType[..., "num_bins"], TensorType[..., "num_bins"]]: (x_start, x_end) - """ + # Returns: + # Tuple[TensorType[..., "num_bins"], TensorType[..., "num_bins"]]: (x_start, x_end) + # """ - x_start = torch.max(torch.where(mask, x[..., None], x[..., :1, None]), -2)[0] - x_end = torch.min(torch.where(~mask, x[..., None], x[..., -1:, None]), -2)[0] - return x_start, x_end + # x_start = torch.max(torch.where(mask, x[..., None], x[..., :1, None]), -2)[0] + # x_end = torch.min(torch.where(~mask, x[..., None], x[..., -1:, None]), -2)[0] + # return x_start, x_end # Force bins to not have a gap between them. Kinda hacky, should reconsider. existing_bins = torch.cat( @@ -344,8 +344,20 @@ class PDFSampler(Sampler): axis=-1, ) - bins_g0, bins_g1 = find_interval(mask, existing_bins) - cdf_g0, cdf_g1 = find_interval(mask, cdf) + # _bins_g0, _bins_g1 = find_interval(mask, existing_bins) + # _cdf_g0, _cdf_g1 = find_interval(mask, cdf) + + # ~7x faster than `find_interval` + u = u.contiguous() + inds = torch.searchsorted(cdf, u, side="right") + below = torch.clamp(inds - 1, 0, num_bins - 1) + above = torch.clamp(inds, 0, num_bins - 1) + cdf_g0 = torch.gather(cdf, -1, below) + bins_g0 = torch.gather(existing_bins, -1, below) + cdf_g1 = torch.gather(cdf, -1, above) + bins_g1 = torch.gather(existing_bins, -1, above) + # assert torch.isclose(_bins_g0, bins_g0).all() + # assert torch.isclose(_cdf_g1, cdf_g1).all() t = torch.clip(torch.nan_to_num((u - cdf_g0) / (cdf_g1 - cdf_g0), 0), 0, 1) bins = bins_g0 + t * (bins_g1 - bins_g0) From 454556c06117c05d8c89f9c07f055484ee50c94f Mon Sep 17 00:00:00 2001 From: Ruilong Li Date: Tue, 12 Jul 2022 07:12:40 +0000 Subject: [PATCH 2/2] update ray_sampler: remove old code and fix lint --- pyrad/graphs/modules/ray_sampler.py | 39 ++--------------------------- 1 file changed, 2 insertions(+), 37 deletions(-) diff --git a/pyrad/graphs/modules/ray_sampler.py b/pyrad/graphs/modules/ray_sampler.py index db8ccb86b..69f59a331 100644 --- a/pyrad/graphs/modules/ray_sampler.py +++ b/pyrad/graphs/modules/ray_sampler.py @@ -17,7 +17,7 @@ Collection of sampling strategies """ from abc import abstractmethod -from typing import Callable, Optional, Tuple +from typing import Callable, Optional import torch from torch import nn @@ -304,35 +304,7 @@ class PDFSampler(Sampler): # Uniform samples between 0 and 1 u = torch.linspace(0.0, 1.0, steps=num_bins, device=cdf.device) u = u.expand(size=(*cdf.shape[:-1], num_bins)) - - # mask = u[..., None, :] >= cdf[..., :, None] # [num_samples, num_orig_bins, num_bins] - - # # Uses same interval trick as mip-NeRF - # def find_interval( - # mask: TensorType[..., "num_orig_bins", "num_bins"], x: TensorType[..., "num_orig_bins"] - # ) -> Tuple[TensorType[..., "num_bins"], TensorType[..., "num_bins"]]: - # """Find intervals based on cdf mask. - - # Mask x x_start x_end - # T T T T F [x0 x1 x2 x3 x4] [x0 x2 x3 x3] [x1 x3 x4 x4] - # T T T T F - # T T T F F - # T F F F F - - # Where the number of rows correspond to the target number of bins. The number of columns - # correspond to the input number of bins - - # Args: - # mask (TensorType[..., "num_orig_bins", "num_bins"]): PDF represented as a boolean mask. - # x (TensorType[..., "num_original_bins"]): Probe to calculate intervals for. - - # Returns: - # Tuple[TensorType[..., "num_bins"], TensorType[..., "num_bins"]]: (x_start, x_end) - # """ - - # x_start = torch.max(torch.where(mask, x[..., None], x[..., :1, None]), -2)[0] - # x_end = torch.min(torch.where(~mask, x[..., None], x[..., -1:, None]), -2)[0] - # return x_start, x_end + u = u.contiguous() # Force bins to not have a gap between them. Kinda hacky, should reconsider. existing_bins = torch.cat( @@ -344,11 +316,6 @@ class PDFSampler(Sampler): axis=-1, ) - # _bins_g0, _bins_g1 = find_interval(mask, existing_bins) - # _cdf_g0, _cdf_g1 = find_interval(mask, cdf) - - # ~7x faster than `find_interval` - u = u.contiguous() inds = torch.searchsorted(cdf, u, side="right") below = torch.clamp(inds - 1, 0, num_bins - 1) above = torch.clamp(inds, 0, num_bins - 1) @@ -356,8 +323,6 @@ class PDFSampler(Sampler): bins_g0 = torch.gather(existing_bins, -1, below) cdf_g1 = torch.gather(cdf, -1, above) bins_g1 = torch.gather(existing_bins, -1, above) - # assert torch.isclose(_bins_g0, bins_g0).all() - # assert torch.isclose(_cdf_g1, cdf_g1).all() t = torch.clip(torch.nan_to_num((u - cdf_g0) / (cdf_g1 - cdf_g0), 0), 0, 1) bins = bins_g0 + t * (bins_g1 - bins_g0)