From 5da146e65f44f7fcf4c70f9019ed43247087cb4a Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Mon, 18 Jul 2022 09:07:34 -0700 Subject: [PATCH] Faster instant-ngp (#166) --- pyrad/graphs/instant_ngp.py | 37 ++++++++----------------------------- 1 file changed, 8 insertions(+), 29 deletions(-) diff --git a/pyrad/graphs/instant_ngp.py b/pyrad/graphs/instant_ngp.py index 855147232..78a280971 100644 --- a/pyrad/graphs/instant_ngp.py +++ b/pyrad/graphs/instant_ngp.py @@ -30,7 +30,7 @@ from pyrad.graphs.base import Graph from pyrad.optimizers.loss import MSELoss from pyrad.fields.occupancy_fields.occupancy_grid import OccupancyGrid from pyrad.renderers.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer -from pyrad.graphs.modules.ray_sampler import PDFSampler, UniformSampler +from pyrad.graphs.modules.ray_sampler import UniformSampler from pyrad.utils import colors from pyrad.cameras.rays import RayBundle from pyrad.utils import visualization, writer @@ -67,10 +67,10 @@ class NGPGraph(Graph): self.occupancy_grid = OccupancyGrid(aabb=self.scene_bounds.aabb) # samplers - self.sampler_occupancy_grid = UniformSampler(num_samples=128, occupancy_field=self.occupancy_grid) - # NOTE(ethan): are we sure we want the include_original flag used like this? - # it could be easily forgotten that it's by default True...? - self.sampler_pdf = PDFSampler(num_samples=128, include_original=False) + self.sampler_occupancy_grid = UniformSampler(num_samples=128) + + # TODO stabalize occupancy grid. + # self.sampler_occupancy_grid = UniformSampler(num_samples=128, occupancy_field=self.occupancy_grid) # renderers self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) @@ -93,30 +93,9 @@ class NGPGraph(Graph): def get_outputs(self, ray_bundle: RayBundle): # uniform sampling - ray_samples_uniform = self.sampler_occupancy_grid(ray_bundle) - field_outputs_uniform = self.field.forward(ray_samples_uniform) - weights_uniform = ray_samples_uniform.get_weights(field_outputs_uniform[FieldHeadNames.DENSITY]) - - # pdf sampling - ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_uniform) - field_outputs_pdf = self.field.forward(ray_samples_pdf) - - # Hacky treatment of bins as points to allow us to merge uniform and pdf. - ts_uniform = (ray_samples_uniform.frustums.starts + ray_samples_uniform.frustums.ends) / 2.0 - ts_pdf = (ray_samples_pdf.frustums.starts + ray_samples_pdf.frustums.ends) / 2.0 - ts, indices = torch.sort(torch.cat([ts_uniform, ts_pdf], -2), -2) - bin_starts = ts - bin_ends = torch.cat([ts[..., 1:, :], ts[..., -1:, :]], dim=-2) - ray_samples = ray_bundle.get_ray_samples(bin_starts=bin_starts, bin_ends=bin_ends) - - field_outputs = {} - for fo_name, _ in field_outputs_pdf.items(): - fo_uniform = field_outputs_uniform[fo_name] - fo_pdf = field_outputs_pdf[fo_name] - fo_uniform_pdf = torch.cat([fo_uniform, fo_pdf], 1) - index = indices.view(fo_uniform_pdf[..., :1].shape) - index = index.expand(-1, -1, fo_uniform_pdf.shape[-1]) # TODO: don't hardcode this - field_outputs[fo_name] = torch.gather(fo_uniform_pdf, dim=1, index=index) + ray_samples = self.sampler_occupancy_grid(ray_bundle) + field_outputs = self.field.forward(ray_samples) + weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY]) weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY]) rgb = self.renderer_rgb(