From a04427644a28e35a8a544b89bc5eb47dbef38498 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Tue, 20 Sep 2022 15:10:48 -0700 Subject: [PATCH] Add Adam optimizer (#496) * Add adam optimizer * Fix configs * merge --- nerfactory/configs/base.py | 24 ++++++++-- nerfactory/configs/base_configs.py | 75 +++++++++++++++++++++--------- 2 files changed, 73 insertions(+), 26 deletions(-) diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index e58d8ec3d..8a9ce33d6 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -337,18 +337,34 @@ class ViewerConfig(PrintableConfig): # Optimizer related configs @dataclass -class OptimizerConfig(InstantiateConfig): +class OptimizerConfig(PrintableConfig): """Basic optimizer config with RAdam""" - _target: Type = torch.optim.RAdam + _target: Type = torch.optim.Adam lr: float = 0.0005 eps: float = 1e-08 # 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, params=None, **kwargs) -> Any: + def setup(self, params) -> Any: """Returns the instantiated object using the config.""" - return self._target(params, lr=self.lr, eps=self.eps) + kwargs = vars(self) + del kwargs["_target"] + return self._target(params, **kwargs) + + +@dataclass +class AdamOptimizerConfig(OptimizerConfig): + """Basic optimizer config with Adam""" + + _target: Type = torch.optim.Adam + + +@dataclass +class RAdamOptimizerConfig(OptimizerConfig): + """Basic optimizer config with RAdam""" + + _target: Type = torch.optim.RAdam @dataclass diff --git a/nerfactory/configs/base_configs.py b/nerfactory/configs/base_configs.py index ea7fc7c27..f517b0aca 100644 --- a/nerfactory/configs/base_configs.py +++ b/nerfactory/configs/base_configs.py @@ -23,13 +23,14 @@ import dcargs from typeguard import typeguard_ignore from nerfactory.configs.base import ( + AdamOptimizerConfig, BlenderDataParserConfig, Config, LoggingConfig, MipNerf360DataParserConfig, NerfWModelConfig, - OptimizerConfig, PipelineConfig, + RAdamOptimizerConfig, SchedulerConfig, TensoRFModelConfig, TrainerConfig, @@ -46,6 +47,23 @@ from nerfactory.models.semantic_nerf import SemanticNerfModel from nerfactory.models.vanilla_nerf import NeRFModel base_configs: Dict[str, Config] = {} +base_configs["compound"] = Config( + method_name="compound", + trainer=TrainerConfig(mixed_precision=True), + pipeline=PipelineConfig( + datamanager=VanillaDataManagerConfig(train_dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192), + model=CompoundModelConfig(eval_num_rays_per_chunk=8192), + ), + optimizers={ + "fields": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + } + }, + viewer=ViewerConfig(enable=True), + logging=LoggingConfig(event_writer="none"), +) + base_configs["instant-ngp"] = Config( method_name="instant-ngp", trainer=TrainerConfig(steps_per_eval_batch=500, steps_per_save=2000, mixed_precision=True), @@ -55,7 +73,7 @@ base_configs["instant-ngp"] = Config( ), optimizers={ "fields": { - "optimizer": OptimizerConfig(lr=3e-3, eps=1e-15), + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), "scheduler": None, } }, @@ -79,6 +97,12 @@ base_configs["mipnerf-360"] = Config( eval_num_rays_per_chunk=8192, ), ), + optimizers={ + "fields": { + "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), + "scheduler": None, + } + }, ) base_configs["mipnerf"] = Config( @@ -93,6 +117,12 @@ base_configs["mipnerf"] = Config( eval_num_rays_per_chunk=8192, ), ), + optimizers={ + "fields": { + "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), + "scheduler": None, + } + }, ) base_configs["nerfw"] = Config( @@ -103,6 +133,12 @@ base_configs["nerfw"] = Config( ), model=NerfWModelConfig(), ), + optimizers={ + "fields": { + "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), + "scheduler": None, + } + }, ) @@ -119,6 +155,12 @@ base_configs["semantic-nerf"] = Config( num_importance_samples=64, ), ), + optimizers={ + "fields": { + "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), + "scheduler": None, + } + }, ) base_configs["vanilla-nerf"] = Config( @@ -129,6 +171,12 @@ base_configs["vanilla-nerf"] = Config( ), model=VanillaModelConfig(_target=NeRFModel), ), + optimizers={ + "fields": { + "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), + "scheduler": None, + } + }, ) base_configs["tensorf"] = Config( @@ -142,37 +190,20 @@ base_configs["tensorf"] = Config( ), optimizers={ "fields": { - "optimizer": OptimizerConfig(lr=0.001), + "optimizer": RAdamOptimizerConfig(lr=0.001), "scheduler": SchedulerConfig(lr_final=0.00005, max_steps=15000), }, "position_encoding": { - "optimizer": OptimizerConfig(lr=0.02), + "optimizer": RAdamOptimizerConfig(lr=0.02), "scheduler": SchedulerConfig(lr_final=0.005, max_steps=15000), }, "direction_encoding": { - "optimizer": OptimizerConfig(lr=0.02), + "optimizer": RAdamOptimizerConfig(lr=0.02), "scheduler": SchedulerConfig(lr_final=0.005, max_steps=15000), }, }, ) -base_configs["compound"] = Config( - method_name="compound", - trainer=TrainerConfig(mixed_precision=True), - pipeline=PipelineConfig( - datamanager=VanillaDataManagerConfig(train_dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192), - model=CompoundModelConfig(eval_num_rays_per_chunk=8192), - ), - optimizers={ - "fields": { - "optimizer": OptimizerConfig(lr=3e-3, eps=1e-15), - "scheduler": None, - } - }, - viewer=ViewerConfig(enable=True), - logging=LoggingConfig(event_writer="none"), -) - @typeguard_ignore # TypeGuard doesn't understand the generic alias that's returned here. def _make_base_config_subcommand_type() -> Type[Config]: