| """训练相关：多任务损失合并、Trainer、调度器。""" | |
| from .multitask import GradNormBalancer, PCGradCombiner, MultiTaskOptimizer | |
| from .schedule import build_scheduler | |
| from .trainer import Trainer, TrainerConfig, compute_all_losses, MAIN_TASK_KEYS, AUX_TASK_KEYS | |
| __all__ = [ | |
| "GradNormBalancer", | |
| "PCGradCombiner", | |
| "MultiTaskOptimizer", | |
| "build_scheduler", | |
| "Trainer", | |
| "TrainerConfig", | |
| "compute_all_losses", | |
| "MAIN_TASK_KEYS", | |
| "AUX_TASK_KEYS", | |
| ] | |