tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

Refactor schedulers (#1466)

Browse Source 
* Refactor schedulers

* lint

* lint
This commit is contained in:
Matthew Tancik 2023-02-21 12:51:51 -08:00 committed by GitHub
parent 53e99c396b
commit f3a34bbf95
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 68 additions and 79 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
nerfstudio/cameras/camera_optimizers.py
View File

@ -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"

6
nerfstudio/configs/method_configs.py
View File

@ -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),

6
nerfstudio/engine/optimizers.py
View File

@ -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:

118
nerfstudio/engine/schedulers.py
View File

@ -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

10
nerfstudio/models/tensorf.py
View File

@ -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(