37 lines
777 B
Python
37 lines
777 B
Python
"""Optimizer implementations and factory registration."""
|
|
|
|
from torch.optim import Optimizer
|
|
|
|
from astrai.factory import BaseFactory
|
|
|
|
|
|
class OptimizerFactory(BaseFactory[Optimizer]):
|
|
"""Factory for built-in training optimizers."""
|
|
|
|
|
|
from astrai.optim.muon_mix import MuonMix
|
|
from astrai.optim.nora_nadamw import (
|
|
NAdamW,
|
|
Nora,
|
|
NoraNAdamW,
|
|
OptimizerParameterGroups,
|
|
nora_direction,
|
|
nora_lr_scale,
|
|
partition_optimizer_parameters,
|
|
)
|
|
|
|
OptimizerFactory.register("nora_nadamw")(NoraNAdamW)
|
|
OptimizerFactory.register("muon_adamw")(MuonMix)
|
|
|
|
__all__ = [
|
|
"MuonMix",
|
|
"NAdamW",
|
|
"Nora",
|
|
"NoraNAdamW",
|
|
"OptimizerFactory",
|
|
"OptimizerParameterGroups",
|
|
"nora_direction",
|
|
"nora_lr_scale",
|
|
"partition_optimizer_parameters",
|
|
]
|