From 8093a359e82e9799f7d61e62e938fde3537f1cc8 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Mon, 3 Oct 2022 18:38:36 -0700 Subject: [PATCH] Tune proposal network (#638) --- nerfstudio/configs/model_configs.py | 4 +-- nerfstudio/model_components/ray_samplers.py | 16 +++++++++-- nerfstudio/model_components/renderers.py | 29 ++++++++++++------- nerfstudio/models/instant_ngp.py | 2 +- nerfstudio/models/nerfacto.py | 32 +++++++++------------ nerfstudio/viewer/server/viewer_utils.py | 6 +++- 6 files changed, 53 insertions(+), 36 deletions(-) diff --git a/nerfstudio/configs/model_configs.py b/nerfstudio/configs/model_configs.py index 2a3d09bd5..9b1ed8fe4 100644 --- a/nerfstudio/configs/model_configs.py +++ b/nerfstudio/configs/model_configs.py @@ -62,7 +62,7 @@ model_configs["nerfacto"] = Config( datamanager=VanillaDataManagerConfig( dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=8192 ), - model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 16), + model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 14), ), optimizers={ "proposal_networks": { @@ -74,7 +74,7 @@ model_configs["nerfacto"] = Config( "scheduler": None, }, }, - viewer=ViewerConfig(num_rays_per_chunk=1 << 16), + viewer=ViewerConfig(num_rays_per_chunk=1 << 14), vis="viewer", ) diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index 8490d8ac2..2705b516c 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -59,6 +59,7 @@ class SpacedSampler(Sampler): spacing_fn: Function that dictates sample spacing (ie `lambda x : x` is uniform). spacing_fn_inv: The inverse of spacing_fn. train_stratified: Use stratified sampling during training. Defults to True + single_jitter: Use a same random jitter for all samples along a ray. Defaults to False """ def __init__( @@ -131,6 +132,7 @@ class UniformSampler(SpacedSampler): Args: num_samples: Number of samples per ray train_stratified: Use stratified sampling during training. Defults to True + single_jitter: Use a same random jitter for all samples along a ray. Defaults to False """ def __init__( @@ -154,6 +156,7 @@ class LinearDisparitySampler(SpacedSampler): Args: num_samples: Number of samples per ray train_stratified: Use stratified sampling during training. Defults to True + single_jitter: Use a same random jitter for all samples along a ray. Defaults to False """ def __init__( @@ -225,6 +228,7 @@ class UniformLinDispPiecewiseSampler(SpacedSampler): Args: num_samples: Number of samples per ray train_stratified: Use stratified sampling during training. Defults to True + single_jitter: Use a same random jitter for all samples along a ray. Defaults to False """ def __init__( @@ -248,6 +252,7 @@ class PDFSampler(Sampler): Args: num_samples: Number of samples per ray train_stratified: Randomize location within each bin during training. + single_jitter: Use a same random jitter for all samples along a ray. Defaults to False include_original: Add original samples to ray. histogram_padding: Amount to weights prior to computing PDF. """ @@ -256,6 +261,7 @@ class PDFSampler(Sampler): self, num_samples: Optional[int] = None, train_stratified: bool = True, + single_jitter: bool = False, include_original: bool = True, histogram_padding: float = 0.01, ) -> None: @@ -263,6 +269,7 @@ class PDFSampler(Sampler): self.train_stratified = train_stratified self.include_original = include_original self.histogram_padding = histogram_padding + self.single_jitter = single_jitter def generate_ray_samples( self, @@ -308,7 +315,12 @@ class PDFSampler(Sampler): # Stratified samples between 0 and 1 u = torch.linspace(0.0, 1.0 - (1.0 / num_bins), steps=num_bins, device=cdf.device) u = u.expand(size=(*cdf.shape[:-1], num_bins)) - u = u + torch.rand(size=(*cdf.shape[:-1], num_bins), device=cdf.device) / num_bins + # u = u + torch.rand(size=(*cdf.shape[:-1], num_bins), device=cdf.device) / num_bins + if self.single_jitter: + rand = torch.rand((*cdf.shape[:-1], 1), device=cdf.device) / num_bins + else: + rand = torch.rand((*cdf.shape[:-1], num_samples + 1), device=cdf.device) / num_bins + u = u + rand else: # Uniform samples between 0 and 1 u = torch.linspace(0.0, 1.0 - (1.0 / num_bins), steps=num_bins, device=cdf.device) @@ -505,7 +517,7 @@ class ProposalNetworkSampler(Sampler): # samplers self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter) - self.pdf_sampler = PDFSampler(include_original=False) + self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter) self._anneal = 1.0 diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index c85004ea1..1ebb61015 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -27,7 +27,7 @@ Example: """ import math -from typing import Literal, Optional +from typing import Literal, Optional, Union import nerfacc import torch @@ -45,7 +45,7 @@ class RGBRenderer(nn.Module): background_color: Background color as RGB. Uses random colors if None. """ - def __init__(self, background_color: Optional[TensorType[3]] = None) -> None: + def __init__(self, background_color: Union[Literal["random", "last_sample"], TensorType[3]] = "random") -> None: super().__init__() self.background_color = background_color @@ -54,7 +54,7 @@ class RGBRenderer(nn.Module): cls, rgb: TensorType["bs":..., "num_samples", 3], weights: TensorType["bs":..., "num_samples", 1], - background_color: Optional[TensorType[3]] = None, + background_color: Union[Literal["random", "last_sample"], TensorType[3]] = "random", ray_indices: Optional[TensorType["num_samples"]] = None, num_rays: Optional[int] = None, ) -> TensorType["bs":..., 3]: @@ -63,7 +63,7 @@ class RGBRenderer(nn.Module): Args: rgb: RGB for each sample weights: Weights for each sample - background_color: Background color as RGB. Uses random colors if None. + background_color: Background color as RGB. ray_indices: Ray index for each sample, used when samples are packed. num_rays: Number of rays, used when samples are packed. @@ -72,18 +72,23 @@ class RGBRenderer(nn.Module): """ if ray_indices is not None and num_rays is not None: # Necessary for packed samples from volumetric ray sampler - rgb = nerfacc.accumulate_along_rays(weights, ray_indices, rgb, num_rays) + if background_color == "last_sample": + raise NotImplementedError("Background color 'last_sample' not implemented for packed samples.") + comp_rgb = nerfacc.accumulate_along_rays(weights, ray_indices, rgb, num_rays) accumulated_weight = nerfacc.accumulate_along_rays(weights, ray_indices, None, num_rays) else: - rgb = torch.sum(weights * rgb, dim=-2) + comp_rgb = torch.sum(weights * rgb, dim=-2) accumulated_weight = torch.sum(weights, dim=-2) - if background_color is None: - background_color = torch.rand_like(rgb).to(rgb.device) + if background_color == "last_sample": + background_color = rgb[..., -1, :] + if background_color == "random": + background_color = torch.rand_like(comp_rgb).to(rgb.device) - rgb = rgb + background_color.to(weights.device) * (1.0 - accumulated_weight) + assert isinstance(background_color, torch.Tensor) + comp_rgb = comp_rgb + background_color.to(weights.device) * (1.0 - accumulated_weight) - return rgb + return comp_rgb def forward( self, @@ -121,7 +126,9 @@ class SHRenderer(nn.Module): """ def __init__( - self, background_color: Optional[TensorType[3]] = None, activation: Optional[nn.Module] = nn.Sigmoid() + self, + background_color: Union[Literal["random", "last_sample"], TensorType[3]] = "random", + activation: Optional[nn.Module] = nn.Sigmoid(), ) -> None: super().__init__() self.background_color = background_color diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index 12abb4ac7..7df255921 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/nerfstudio/models/instant_ngp.py @@ -122,7 +122,7 @@ class NGPModel(Model): ) # renderers - background_color = None if self.config.randomize_background else colors.WHITE + background_color = "random" if self.config.randomize_background else colors.WHITE self.renderer_rgb = RGBRenderer(background_color=background_color) self.renderer_accumulation = AccumulationRenderer() self.renderer_depth = DepthRenderer() diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 3f094a1c8..11c4088e9 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -19,7 +19,7 @@ NeRF implementation that combines many recent advancements. from __future__ import annotations from dataclasses import dataclass, field -from typing import Dict, List, Tuple, Type +from typing import Dict, List, Literal, Tuple, Type import numpy as np import torch @@ -47,7 +47,7 @@ from nerfstudio.model_components.renderers import ( ) from nerfstudio.model_components.scene_colliders import NearFarCollider from nerfstudio.models.base_model import Model, ModelConfig -from nerfstudio.utils import colormaps, colors +from nerfstudio.utils import colormaps @dataclass @@ -59,7 +59,7 @@ class NerfactoModelConfig(ModelConfig): """How far along the ray to start sampling.""" far_plane: float = 1000.0 """How far along the ray to stop sampling.""" - randomize_background: bool = True + background_color: Literal["background", "last_sample"] = "last_sample" """Whether to randomize the background color.""" num_proposal_samples_per_ray: Tuple[int] = (64,) """Number of samples per ray for the proposal network.""" @@ -71,10 +71,8 @@ class NerfactoModelConfig(ModelConfig): """Use the same proposal network. Otherwise use different ones.""" interlevel_loss_mult: float = 1.0 """Proposal loss multiplier.""" - distortion_loss_mult: float = 0.01 + distortion_loss_mult: float = 0.002 """Distortion loss multiplier.""" - use_appearance_conditioning: bool = True - """Whether to use appearance conditioning.""" use_proposal_weight_anneal: bool = True """Whether to use proposal weight annealing.""" use_average_appearance_embedding: bool = True @@ -83,8 +81,8 @@ class NerfactoModelConfig(ModelConfig): """Slope of the annealing function for the proposal weights.""" proposal_weights_anneal_max_num_iters: int = 1000 """Max num iterations for the annealing function.""" - use_single_jitter: bool = False - """Whether use single jitter or not for first proposal network.""" + use_single_jitter: bool = True + """Whether use single jitter or not for the proposal networks.""" class NerfactoModel(Model): @@ -103,15 +101,12 @@ class NerfactoModel(Model): scene_contraction = SceneContraction(order=float("inf")) # Fields - if self.config.use_appearance_conditioning: - self.field = TCNNNerfactoField( - self.scene_box.aabb, - spatial_distortion=scene_contraction, - num_images=self.num_train_data, - use_average_appearance_embedding=self.config.use_average_appearance_embedding, - ) - else: - raise NotImplementedError("Only appearance conditioning is supported.") + self.field = TCNNNerfactoField( + self.scene_box.aabb, + spatial_distortion=scene_contraction, + num_images=self.num_train_data, + use_average_appearance_embedding=self.config.use_average_appearance_embedding, + ) # Build the proposal network(s) self.proposal_networks = torch.nn.ModuleList() @@ -137,8 +132,7 @@ class NerfactoModel(Model): ) # renderers - background_color = None if self.config.randomize_background else colors.WHITE - self.renderer_rgb = RGBRenderer(background_color=background_color) + self.renderer_rgb = RGBRenderer(background_color=self.config.background_color) self.renderer_accumulation = AccumulationRenderer() self.renderer_depth = DepthRenderer() diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 8def39962..e975dffd3 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -537,8 +537,12 @@ class ViewerState: image_height = (num_vis_rays / aspect_ratio) ** 0.5 image_height = int(round(image_height, -1)) image_height = min(self.max_resolution, image_height) + image_width = int(image_height * aspect_ratio) + if image_width > self.max_resolution: + image_width = self.max_resolution + image_height = int(image_width / aspect_ratio) - return image_height, int(image_height * aspect_ratio) + return image_height, image_width def _process_invalid_output(self, output_type: str) -> str: """Check to see whether we are in the corner case of RGB; if still invalid, throw error