From 2150bad0e4ca14339f55fa620c209dc3e00d2b10 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Tue, 11 Oct 2022 14:21:00 -0700 Subject: [PATCH] Reducing flicker artifacts for nerfacto method (#750) * pose opt default nerfacto. * added SO3 cam pose optimization which is slightly faster (~2.2m/iter) than BARF * added occupancy grid with placeholder until nerfacc is updated to query density from it * added occupancy grid sampling as the level 0 in proposal hierarchy, relies on nerfacc 0.1.7 (on pypi) * merge main * merge main * minor changes to sampling * adjusted sample number * refactor occ grid stuff, settle on decent params * nograd proposal sampling * param tweaks * varying resolution proposal networks * params * merged main * optimizer update * remove occ grid * remove double nerfacto samples * set single jitter to true * remove pdb * more changes * fix method defaults * fix pylint * remove occupancy fn Co-authored-by: akristoffersen Co-authored-by: Justin Kerr --- nerfstudio/configs/method_configs.py | 6 ++-- nerfstudio/engine/optimizers.py | 1 + nerfstudio/fields/density_fields.py | 10 +++--- nerfstudio/model_components/ray_samplers.py | 4 +-- nerfstudio/models/nerfacto.py | 40 ++++++++++++++------- nerfstudio/utils/eval_utils.py | 12 +------ pyproject.toml | 1 + scripts/render.py | 3 -- 8 files changed, 40 insertions(+), 37 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 1dd473a8a..69fc4a59a 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -56,8 +56,10 @@ method_configs["nerfacto"] = Config( datamanager=VanillaDataManagerConfig( dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, - eval_num_rays_per_batch=8192, - camera_optimizer=CameraOptimizerConfig(mode="SO3xR3"), + eval_num_rays_per_batch=4096, + camera_optimizer=CameraOptimizerConfig( + mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-3) + ), ), model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 14), ), diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py index 500e2328a..5bae109d0 100644 --- a/nerfstudio/engine/optimizers.py +++ b/nerfstudio/engine/optimizers.py @@ -51,6 +51,7 @@ class AdamOptimizerConfig(OptimizerConfig): """Basic optimizer config with Adam""" _target: Type = torch.optim.Adam + weight_decay: float = 0 @dataclass diff --git a/nerfstudio/fields/density_fields.py b/nerfstudio/fields/density_fields.py index 20dc35163..af7b534ed 100644 --- a/nerfstudio/fields/density_fields.py +++ b/nerfstudio/fields/density_fields.py @@ -55,6 +55,11 @@ class HashMLPDensityField(Field): hidden_dim: int = 64, spatial_distortion: Optional[SpatialDistortion] = None, use_linear=False, + num_levels=8, + max_res=1024, + base_res=16, + log2_hashmap_size=18, + features_per_level=2, ) -> None: super().__init__() @@ -62,11 +67,6 @@ class HashMLPDensityField(Field): self.spatial_distortion = spatial_distortion self.use_linear = use_linear - num_levels = 8 - max_res = 1024 - base_res = 16 - log2_hashmap_size = 18 - features_per_level = 2 growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1)) config = { diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index 1f3dac2ea..f744e66ae 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -315,7 +315,6 @@ 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 if self.single_jitter: rand = torch.rand((*cdf.shape[:-1], 1), device=cdf.device) / num_bins else: @@ -469,7 +468,6 @@ class VolumetricSampler(Sampler): cone_angle=cone_angle, alpha_thre=1e-2, ) - num_samples = starts.shape[0] if num_samples == 0: # create a single fake sample and update packed_info accordingly @@ -506,7 +504,7 @@ class ProposalNetworkSampler(Sampler): num_proposal_samples_per_ray: Tuple[int] = (64,), num_nerf_samples_per_ray: int = 32, num_proposal_network_iterations: int = 2, - single_jitter: bool = True, + single_jitter: bool = False, ) -> None: super().__init__() self.num_proposal_samples_per_ray = num_proposal_samples_per_ray diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 11c4088e9..b368d1cb4 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -61,14 +61,21 @@ class NerfactoModelConfig(ModelConfig): """How far along the ray to stop sampling.""" background_color: Literal["background", "last_sample"] = "last_sample" """Whether to randomize the background color.""" - num_proposal_samples_per_ray: Tuple[int] = (64,) + num_proposal_samples_per_ray: Tuple[int] = (200, 64) """Number of samples per ray for the proposal network.""" num_nerf_samples_per_ray: int = 64 """Number of samples per ray for the nerf network.""" - num_proposal_network_iterations: int = 1 + num_proposal_iterations: int = 2 """Number of proposal network iterations.""" use_same_proposal_network: bool = False """Use the same proposal network. Otherwise use different ones.""" + proposal_net_args_list: List[Dict] = field( + default_factory=lambda: [ + {"hidden_dim": 16, "log2_hashmap_size": 16, "num_levels": 5, "max_res": 64}, + {"hidden_dim": 16, "log2_hashmap_size": 16, "num_levels": 5, "max_res": 256}, + ] + ) + """Arguments for the proposal density fields.""" interlevel_loss_mult: float = 1.0 """Proposal loss multiplier.""" distortion_loss_mult: float = 0.002 @@ -108,29 +115,36 @@ class NerfactoModel(Model): use_average_appearance_embedding=self.config.use_average_appearance_embedding, ) + self.density_fns = [] + num_prop_nets = self.config.num_proposal_iterations # Build the proposal network(s) self.proposal_networks = torch.nn.ModuleList() if self.config.use_same_proposal_network: - network = HashMLPDensityField(self.scene_box.aabb, spatial_distortion=scene_contraction) + assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed." + prop_net_args = self.config.proposal_net_args_list[0] + network = HashMLPDensityField(self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args) self.proposal_networks.append(network) - self.density_fns = [network.density_fn for _ in range(self.config.num_proposal_network_iterations)] + self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)]) else: - for _ in range(self.config.num_proposal_network_iterations): - network = HashMLPDensityField(self.scene_box.aabb, spatial_distortion=scene_contraction) + for i in range(num_prop_nets): + prop_net_args = self.config.proposal_net_args_list[min(i, len(self.config.proposal_net_args_list) - 1)] + network = HashMLPDensityField( + self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args + ) self.proposal_networks.append(network) - self.density_fns = [network.density_fn for network in self.proposal_networks] - - # Collider - self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane) + self.density_fns.extend([network.density_fn for network in self.proposal_networks]) # Samplers self.proposal_sampler = ProposalNetworkSampler( num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray, num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray, - num_proposal_network_iterations=self.config.num_proposal_network_iterations, + num_proposal_network_iterations=self.config.num_proposal_iterations, single_jitter=self.config.use_single_jitter, ) + # Collider + self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane) + # renderers self.renderer_rgb = RGBRenderer(background_color=self.config.background_color) self.renderer_accumulation = AccumulationRenderer() @@ -189,7 +203,7 @@ class NerfactoModel(Model): outputs["weights_list"] = weights_list outputs["ray_samples_list"] = ray_samples_list - for i in range(self.config.num_proposal_network_iterations): + for i in range(self.config.num_proposal_iterations): outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i]) return outputs @@ -242,7 +256,7 @@ class NerfactoModel(Model): images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth} - for i in range(self.config.num_proposal_network_iterations): + for i in range(self.config.num_proposal_iterations): key = f"prop_depth_{i}" prop_depth_i = colormaps.apply_depth_colormap( outputs[key], diff --git a/nerfstudio/utils/eval_utils.py b/nerfstudio/utils/eval_utils.py index 5e3cf4c3d..a97997ffa 100644 --- a/nerfstudio/utils/eval_utils.py +++ b/nerfstudio/utils/eval_utils.py @@ -63,9 +63,7 @@ def eval_load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path: return load_path -def eval_setup( - config_path: Path, double_nerfacto_nerf_samples: bool = False, eval_num_rays_per_chunk: Optional[int] = None -) -> Tuple[cfg.Config, Pipeline, Path]: +def eval_setup(config_path: Path, eval_num_rays_per_chunk: Optional[int] = None) -> Tuple[cfg.Config, Pipeline, Path]: """Shared setup for loading a saved pipeline for evaluation. Args: @@ -81,14 +79,6 @@ def eval_setup( if eval_num_rays_per_chunk: config.pipeline.model.eval_num_rays_per_chunk = eval_num_rays_per_chunk - if config.method_name == "nerfacto" and double_nerfacto_nerf_samples: - config.pipeline.model.num_nerf_samples_per_ray *= 2 - console.print( - "[bold yellow]Warning: doubling the number of nerfacto nerf samples to " - f"{config.pipeline.model.num_nerf_samples_per_ray} for improved quality." - ) - console.print("[bold yellow]This logic is a hack and will be removed in the future.") - # load checkpoints from wherever they were saved # TODO: expose the ability to choose an arbitrary checkpoint config.trainer.load_dir = config.get_checkpoint_dir() diff --git a/pyproject.toml b/pyproject.toml index e80a7d023..c5ae4fba1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -107,6 +107,7 @@ generated-members = ["numpy.*", "torch.*", "cv2.*", "cv.*"] good-names-rgxs = "^[_a-zA-Z][_a-z0-9]?$" ignore-paths = ["scripts/colmap2nerf.py"] jobs = 0 +ignored-classes = ["TensorDataclass"] disable = [ "duplicate-code", diff --git a/scripts/render.py b/scripts/render.py index e785c2ffd..89b92c8a1 100644 --- a/scripts/render.py +++ b/scripts/render.py @@ -106,8 +106,6 @@ class RenderTrajectory: output_path: Path = Path("renders/output.mp4") # How long the video should be. seconds: float = 5.0 - # A hack to double the number of samples for the nerfacto method. - double_nerfacto_nerf_samples: bool = True # Specifies number of rays per chunk during eval. eval_num_rays_per_chunk: Optional[int] = None @@ -115,7 +113,6 @@ class RenderTrajectory: """Main function.""" _, pipeline, _ = eval_setup( self.load_config, - double_nerfacto_nerf_samples=self.double_nerfacto_nerf_samples, eval_num_rays_per_chunk=self.eval_num_rays_per_chunk, )