mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 02:45:34 +08:00
66 lines
2.2 KiB
Python
66 lines
2.2 KiB
Python
"""
|
|
Optimizers class.
|
|
"""
|
|
|
|
from typing import Dict, List
|
|
from omegaconf import DictConfig
|
|
import torch
|
|
from torch.nn import Parameter
|
|
|
|
|
|
class Optimizers:
|
|
"""_summary_"""
|
|
|
|
def __init__(self, config: DictConfig, param_groups: Dict[str, List[Parameter]]):
|
|
"""_summary_
|
|
|
|
Args:
|
|
config (DictConfig): _description_
|
|
param_dict (Dict[str, List[Parameter]]): _description_
|
|
"""
|
|
self.config = config
|
|
self.optimizers = {}
|
|
self.schedulers = {}
|
|
for param_name, params in param_groups.items():
|
|
optimizer_config = config[param_name].optimizer
|
|
optimizer = getattr(torch.optim, optimizer_config.type)
|
|
kwargs = {k: v for k, v in optimizer_config.items() if k != "type"}
|
|
self.optimizers[param_name] = optimizer(params, **kwargs)
|
|
if config[param_name].scheduler:
|
|
scheduler_config = config[param_name].scheduler
|
|
scheduler = getattr(torch.optim.lr_scheduler, scheduler_config.type)
|
|
kwargs = {k: v for k, v in scheduler_config.items() if k != "type"}
|
|
self.schedulers[param_name] = scheduler(self.optimizers[param_name], **kwargs)
|
|
|
|
def optimizer_step(self, param_name: str) -> None:
|
|
"""fetch and step corresponding optimizer
|
|
|
|
Args:
|
|
param_name (str): name of optimizer to step forward
|
|
"""
|
|
self.optimizers[param_name].step()
|
|
|
|
def scheduler_step(self, param_name: str) -> None:
|
|
"""fetch and step corresponding scheduler
|
|
|
|
Args:
|
|
param_name (str): name of scheduler to step forward
|
|
"""
|
|
if self.config.param_name.scheduler:
|
|
self.schedulers[param_name].step()
|
|
|
|
def zero_grad_all(self):
|
|
"""zero the gradients for all optimizer parameters"""
|
|
for _, optimizer in self.optimizers.items():
|
|
optimizer.zero_grad()
|
|
|
|
def optimizer_step_all(self):
|
|
"""Run step for all optimizers."""
|
|
for _, optimizer in self.optimizers.items():
|
|
# note that they key is the parameter name
|
|
optimizer.step()
|
|
|
|
def scheduler_step_all(self, step):
|
|
"""Run step for all schedulers."""
|
|
# TODO: this needs to be implemented
|