From 5e86d16d742e2c910ecc08f660958106ba2c75f5 Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 21 Feb 2023 17:38:15 +0000 Subject: [PATCH] LR warmup, angle sampling, huge interlevel loss --- nerfstudio/configs/method_configs.py | 20 +++++++++---------- .../datamanagers/dreamfusion_datamanager.py | 9 ++++++--- nerfstudio/engine/schedulers.py | 19 ++++++++++++++++-- nerfstudio/model_components/ray_samplers.py | 2 +- nerfstudio/models/dreamfusion.py | 8 ++++---- 5 files changed, 38 insertions(+), 20 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index a516ba323..448b352a8 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -46,7 +46,7 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import ( PhototourismDataParserConfig, ) from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig -from nerfstudio.engine.schedulers import SchedulerConfig,WarmupScheduler +from nerfstudio.engine.schedulers import SchedulerConfig,WarmupSchedulerConfig from nerfstudio.engine.trainer import TrainerConfig from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig @@ -353,17 +353,17 @@ method_configs["dreamfusion"] = TrainerConfig( model=DreamFusionModelConfig( eval_num_rays_per_chunk=1 << 15, distortion_loss_mult=0.02, - interlevel_loss_mult=1.0, + interlevel_loss_mult=100, orientation_loss_mult=0.1, max_res=256, sphere_collider=True, initialize_density=False, random_background=True, - proposal_warmup=1000, - proposal_update_every=2, + proposal_warmup=500, + proposal_update_every=0, proposal_weights_anneal_max_num_iters=100, - start_lambertian_training=1000, - start_normals_training=500, + start_lambertian_training=500, + start_normals_training=1000, opacity_loss_mult=0.001, prompting_type="location_based", guidance_scale=100, @@ -371,12 +371,12 @@ method_configs["dreamfusion"] = TrainerConfig( ), optimizers={ "proposal_networks": { - "optimizer": AdamOptimizerConfig(lr=5e-2, eps=1e-15), - "scheduler": None, + "optimizer": AdamOptimizerConfig(lr=1e-8, eps=1e-15), + "scheduler": WarmupSchedulerConfig(lr_init=1e-8,lr_max=1e-3,lr_final=1e-6,warmup_steps=2000,max_steps=10000), }, "fields": { - "optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15), - "scheduler": WarmupScheduler(lr_init=1e-9,lr_max=1e-4,lr_final=1e-6,warmup_steps=3000,max_steps=10000), + "optimizer": AdamOptimizerConfig(lr=1e-8, eps=1e-15), + "scheduler": WarmupSchedulerConfig(lr_init=1e-8,lr_max=1e-3,lr_final=1e-6,warmup_steps=2000,max_steps=10000), }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), diff --git a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py index 153913d42..848cfde01 100644 --- a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py +++ b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py @@ -17,7 +17,7 @@ Data manager for dreamfusion """ from __future__ import annotations - +import numpy as np from dataclasses import dataclass, field from typing import Dict, List, Tuple, Type, Union @@ -167,9 +167,9 @@ class DreamFusionDataManagerConfig(VanillaDataManagerConfig): """Mean radius of camera orbit""" radius_std: float = 0.1 """Std of radius of camera orbit""" - focal_range: Tuple[float, float] = (0.6, 1.2) + focal_range: Tuple[float, float] = (0.6, .1.2) """Range of focal length""" - vertical_rotation_range: Tuple[float, float] = (-80, 30) + vertical_rotation_range: Tuple[float, float] = (-90, 0) """Range of vertical rotation""" jitter_std: float = 0.05 """Std of camera direction jitter, so we don't just point the cameras towards the center every time""" @@ -245,6 +245,7 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me # return ray_bundle, {"initialization": False} # TODO below + ang_range = np.interp(step,[0,4000],[0,180]) cameras, vertical_rotation, central_rotation = random_train_pose( self.config.train_images_per_batch, self.config.train_resolution, @@ -255,6 +256,8 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me vertical_rotation_range=self.config.vertical_rotation_range, jitter_std=self.config.jitter_std, center=self.config.center, + central_rotation_range = (-ang_range,ang_range) + ) ray_bundle = cameras.generate_rays(torch.tensor(list(range(self.config.train_images_per_batch)))).flatten() diff --git a/nerfstudio/engine/schedulers.py b/nerfstudio/engine/schedulers.py index 0f9df8da3..4c472bdab 100644 --- a/nerfstudio/engine/schedulers.py +++ b/nerfstudio/engine/schedulers.py @@ -72,6 +72,22 @@ class ExponentialDecaySchedule(lr_scheduler.LambdaLR): super().__init__(optimizer, lr_lambda=func) +@dataclass +class WarmupSchedulerConfig(InstantiateConfig): + """Basic scheduler config with self-defined warmup decay schedule""" + + _target: Type = field(default_factory=lambda: WarmupScheduler) + lr_init:float=1e-8 + lr_max:float=1e-3 + lr_final:float=1e-5 + warmup_steps:int = 1000 + max_steps:int=10000 + # TODO: somehow make this more generic. i dont like the idea of overriding the setup function + # but also not sure how to go about passing things into predefined torch objects. + def setup(self, optimizer=None, **kwargs) -> Any: + """Returns the instantiated object using the config.""" + return self._target(optimizer,self.lr_init, self.lr_max, self.lr_final, self.warmup_steps, self.max_steps) + class WarmupScheduler(lr_scheduler.LambdaLR): """Exponential learning rate decay function. See https://github.com/google-research/google-research/blob/ @@ -90,8 +106,7 @@ class WarmupScheduler(lr_scheduler.LambdaLR): lr = np.interp(step,[0,warmup_steps],[lr_init,lr_max]) else: t = np.clip((step-warmup_steps)/max_steps,0,1) - lr = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t) - print("Lr",lr) + lr = np.exp(np.log(lr_max) * (1 - t) + np.log(lr_final) * t) return lr/lr_init#divide by lr_init to make the lr set to "lr" super().__init__(optimizer, lr_lambda=func) diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index f55795ab5..c2ef4d916 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -263,7 +263,7 @@ class PDFSampler(Sampler): train_stratified: bool = True, single_jitter: bool = False, include_original: bool = True, - histogram_padding: float = 0.00, + histogram_padding: float = 0.005, ) -> None: super().__init__(num_samples=num_samples) self.train_stratified = train_stratified diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 6eba9dc96..fcb76976b 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -114,7 +114,7 @@ class DreamFusionModelConfig(ModelConfig): """Slope of the annealing function for the proposal weights.""" proposal_weights_anneal_max_num_iters: int = 500 """Max num iterations for the annealing function.""" - use_single_jitter: bool = True + use_single_jitter: bool = False """Whether use single jitter or not for the proposal networks.""" interlevel_loss_mult: float = 1.0 """Proposal loss multiplier.""" @@ -126,7 +126,7 @@ class DreamFusionModelConfig(ModelConfig): """start training with lambertian shading after this many iterations""" opacity_penalty: bool = True """enables penalty to encourage sparse weights (penalizing for uniform density along ray)""" - opacity_loss_mult: float = 1 + opacity_loss_mult: float = 1e-3 """scale for opacity penalty""" max_res: int = 256 """Maximum resolution of the density field.""" @@ -341,8 +341,8 @@ class DreamFusionModel(Model): def get_outputs(self, ray_bundle: RayBundle): # uniform sampling - background_rgb = self.field.get_background_rgb(ray_bundle) ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns) + background_rgb = self.field.get_background_rgb(ray_bundle) field_outputs = self.field(ray_samples, compute_normals=True) weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY]) weights_list.append(weights) @@ -435,7 +435,7 @@ class DreamFusionModel(Model): assert weights.shape[-1] == 1 if self.config.opacity_penalty: outputs["opacity_loss"] = ( - torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.opacity_loss_mult + torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) ) return outputs