diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 20aaef3bc..afd002a71 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -31,7 +31,10 @@ from typing_extensions import Literal, assert_never from nerfstudio.cameras.lie_groups import exp_map_SE3, exp_map_SO3xR3 from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.engine.optimizers import AdamOptimizerConfig -from nerfstudio.engine.schedulers import SchedulerConfig +from nerfstudio.engine.schedulers import ( + ExponentialDecaySchedulerConfig, + SchedulerConfig, +) from nerfstudio.utils import poses as pose_utils @@ -53,7 +56,7 @@ class CameraOptimizerConfig(InstantiateConfig): optimizer: AdamOptimizerConfig = AdamOptimizerConfig(lr=6e-4, eps=1e-15) """ADAM parameters for camera optimization.""" - scheduler: SchedulerConfig = SchedulerConfig(max_steps=10000) + scheduler: SchedulerConfig = ExponentialDecaySchedulerConfig(max_steps=10000) """Learning rate scheduler for camera optimizer..""" param_group: tyro.conf.Suppress[str] = "camera_opt" diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 4f41ba4af..6b1a71180 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -43,7 +43,7 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import ( PhototourismDataParserConfig, ) from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig -from nerfstudio.engine.schedulers import SchedulerConfig +from nerfstudio.engine.schedulers import ExponentialDecaySchedulerConfig from nerfstudio.engine.trainer import TrainerConfig from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig @@ -267,11 +267,11 @@ method_configs["tensorf"] = TrainerConfig( optimizers={ "fields": { "optimizer": AdamOptimizerConfig(lr=0.001), - "scheduler": SchedulerConfig(lr_final=0.0001, max_steps=30000), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=30000), }, "encodings": { "optimizer": AdamOptimizerConfig(lr=0.02), - "scheduler": SchedulerConfig(lr_final=0.002, max_steps=30000), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.002, max_steps=30000), }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py index 6878eccc2..786c51115 100644 --- a/nerfstudio/engine/optimizers.py +++ b/nerfstudio/engine/optimizers.py @@ -86,8 +86,10 @@ class Optimizers: self.optimizers[param_group_name] = config[param_group_name]["optimizer"].setup(params=params) self.parameters[param_group_name] = params if config[param_group_name]["scheduler"]: - self.schedulers[param_group_name] = config[param_group_name]["scheduler"].setup( - optimizer=self.optimizers[param_group_name], lr_init=lr_init + self.schedulers[param_group_name] = ( + config[param_group_name]["scheduler"] + .setup() + .get_scheduler(optimizer=self.optimizers[param_group_name], lr_init=lr_init) ) def optimizer_step(self, param_group_name: str) -> None: diff --git a/nerfstudio/engine/schedulers.py b/nerfstudio/engine/schedulers.py index c545b9124..decdcd101 100644 --- a/nerfstudio/engine/schedulers.py +++ b/nerfstudio/engine/schedulers.py @@ -14,8 +14,9 @@ """Scheduler Classes""" +from abc import abstractmethod from dataclasses import dataclass, field -from typing import Any, Optional, Type +from typing import Type import numpy as np from torch.optim import Optimizer, lr_scheduler @@ -25,92 +26,71 @@ from nerfstudio.configs.base_config import InstantiateConfig @dataclass class SchedulerConfig(InstantiateConfig): - """Basic scheduler config with self-defined exponential decay schedule""" + """Basic scheduler config""" - _target: Type = field(default_factory=lambda: ExponentialDecaySchedule) + _target: Type = field(default_factory=lambda: Scheduler) + """_target: target class to instantiate""" + + +class Scheduler: # pylint: disable=too-few-public-methods + """Exponential learning rate decay function.""" + + config: SchedulerConfig + + def __init__(self, config: SchedulerConfig): + super().__init__() + self.config = config + + @abstractmethod + def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler: + """Abstract method that returns a scheduler object. + + Args: + optimizer: The optimizer to use. + lr_init: The initial learning rate. + Returns: + The scheduler object. + """ + + +@dataclass +class ExponentialDecaySchedulerConfig(SchedulerConfig): + """Exponential learning rate decay config""" + + _target: Type = field(default_factory=lambda: ExponentialDecayScheduler) lr_final: float = 0.000005 + """The final learning rate.""" max_steps: int = 1000000 - - # 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, lr_init=None, **kwargs) -> Any: - """Returns the instantiated object using the config.""" - return self._target(optimizer, lr_init, self.lr_final, self.max_steps) + """The maximum number of steps.""" + lr_delay_steps: int = 0 + """The number of steps to delay the learning rate.""" + lr_delay_mult: float = 1.0 + """The multiplier for the learning rate after the delay.""" -class ExponentialDecaySchedule(lr_scheduler.LambdaLR): +class ExponentialDecayScheduler(Scheduler): # pylint: disable=too-few-public-methods """Exponential learning rate decay function. See https://github.com/google-research/google-research/blob/ fd2cea8cdd86b3ed2c640cbe5561707639e682f3/jaxnerf/nerf/utils.py#L360 for details. - - Args: - optimizer: The optimizer to update. - lr_init: The initial learning rate. - lr_final: The final learning rate. - max_steps: The maximum number of steps. - lr_delay_steps: The number of steps to delay the learning rate. - lr_delay_mult: The multiplier for the learning rate after the delay. """ - config: SchedulerConfig + config: ExponentialDecaySchedulerConfig - def __init__(self, optimizer, lr_init, lr_final, max_steps, lr_delay_steps=0, lr_delay_mult=1.0) -> None: + def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler: def func(step): - if lr_delay_steps > 0: - delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin( - 0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1) + if self.config.lr_delay_steps > 0: + delay_rate = self.config.lr_delay_mult + (1 - self.config.lr_delay_mult) * np.sin( + 0.5 * np.pi * np.clip(step / self.config.lr_delay_steps, 0, 1) ) else: delay_rate = 1.0 - t = np.clip(step / max_steps, 0, 1) - log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t) + t = np.clip(step / self.config.max_steps, 0, 1) + log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(self.config.lr_final) * t) multiplier = ( log_lerp / lr_init ) # divided by lr_init because the multiplier is with the initial learning rate return delay_rate * multiplier - super().__init__(optimizer, lr_lambda=func) - - -class DelayerScheduler(lr_scheduler.LambdaLR): - """Starts with a flat lr schedule until it reaches N epochs then applies a given scheduler""" - - def __init__( - self, - optimizer: Optimizer, - lr_init, # pylint: disable=unused-argument - lr_final, # pylint: disable=unused-argument - max_steps, # pylint: disable=unused-argument - delay_epochs: int = 500, - after_scheduler: Optional[lr_scheduler.LambdaLR] = None, - ) -> None: - def func(step): - if step > delay_epochs: - if after_scheduler is not None: - multiplier = after_scheduler.lr_lambdas[0](step - delay_epochs) # type: ignore - return multiplier - return 1.0 - return 0.0 - - super().__init__(optimizer, lr_lambda=func) - - -class DelayedExponentialScheduler(DelayerScheduler): - """Delayer Scheduler with an Exponential Scheduler initialized afterwards.""" - - def __init__( - self, - optimizer: Optimizer, - lr_init, - lr_final, - max_steps, - delay_epochs: int = 200, - ): - after_scheduler = ExponentialDecaySchedule( - optimizer, - lr_init, - lr_final, - max_steps, - ) - super().__init__(optimizer, lr_init, lr_final, max_steps, delay_epochs, after_scheduler=after_scheduler) + scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=func) + return scheduler diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index d5af1579d..48310fde2 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -130,9 +130,13 @@ class TensoRFModel(Model): "optimizer" ].setup(params=enc) if optimizers_config["encodings"]["scheduler"]: - training_callback_attributes.optimizers.schedulers["encodings"] = optimizers_config["encodings"][ - "scheduler" - ].setup(optimizer=training_callback_attributes.optimizers.optimizers["encodings"], lr_init=lr_init) + training_callback_attributes.optimizers.schedulers["encodings"] = ( + optimizers_config["encodings"]["scheduler"] + .setup() + .get_scheduler( + optimizer=training_callback_attributes.optimizers.optimizers["encodings"], lr_init=lr_init + ) + ) callbacks = [ TrainingCallback(