ViperEkura
10ebd7211f
- 2D 参数用 Newton-Schulz 正交化 + Nesterov 动量更新 - 1D 参数用 AdamW 更新 - 支持 lr/momentum/weight_decay/ns_steps 配置
25 lines
580 B
Python
25 lines
580 B
Python
from astrai.trainer.optim import Muon
|
|
from astrai.trainer.schedule import BaseScheduler, SchedulerFactory
|
|
from astrai.trainer.strategy import BaseStrategy, StrategyFactory
|
|
from astrai.trainer.train_callback import (
|
|
CallbackFactory,
|
|
TrainCallback,
|
|
)
|
|
from astrai.trainer.trainer import Trainer
|
|
|
|
__all__ = [
|
|
# Main trainer
|
|
"Trainer",
|
|
# Optimizer
|
|
"Muon",
|
|
# Strategy factory
|
|
"StrategyFactory",
|
|
"BaseStrategy",
|
|
# Scheduler factory
|
|
"SchedulerFactory",
|
|
"BaseScheduler",
|
|
# Callback factory
|
|
"TrainCallback",
|
|
"CallbackFactory",
|
|
]
|