File size: 3,154 Bytes
a311e94 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """Registry of composable training algorithms.
Algorithms are grouped by category and can be combined freely::
>>> from agron_nets import build_algorithm
>>> algos = [build_algorithm("gradient_clip"), build_algorithm("ema")]
"""
from __future__ import annotations
from typing import Any, Dict, List
from .averaging import EMA, SWA
from .base import AlgorithmContext, BaseAlgorithm
from .data import Curriculum, DifficultyCurriculum
from .objectives import DPO, Distillation, distillation_loss, dpo_loss, sequence_log_probs
from .regularization import LabelSmoothing, Mixup, WeightDecaySchedule
from .stability import (
ActivationCheckpointing,
GradientAccumulation,
GradientClip,
GradientNoise,
ZLoss,
)
from .wrappers import SAM, Lookahead
ALGORITHMS: Dict[str, Dict[str, Any]] = {}
def _register(cls) -> None:
ALGORITHMS[cls.name] = {
"name": cls.name,
"class": cls,
"category": getattr(cls, "category", "misc"),
"description": getattr(cls, "description", ""),
}
for _cls in (
LabelSmoothing,
Mixup,
WeightDecaySchedule,
GradientClip,
GradientAccumulation,
ZLoss,
GradientNoise,
ActivationCheckpointing,
SAM,
Lookahead,
EMA,
SWA,
Distillation,
DPO,
Curriculum,
DifficultyCurriculum,
):
_register(_cls)
def algorithm_names(category: str | None = None) -> List[str]:
if category:
return sorted(k for k, v in ALGORITHMS.items() if v["category"] == category)
return sorted(ALGORITHMS.keys())
def get_algorithm(name: str) -> Dict[str, Any]:
key = (name or "").strip().lower()
if key not in ALGORITHMS:
raise KeyError(f"unknown algorithm '{name}'. available: {', '.join(algorithm_names())}")
return ALGORITHMS[key]
def build_algorithm(name: str, **kwargs) -> BaseAlgorithm:
"""Instantiate one algorithm by registry name."""
return get_algorithm(name)["class"](**kwargs)
def build_algorithms(names, **kwargs) -> List[BaseAlgorithm]:
"""Instantiate several algorithms (composable)."""
if isinstance(names, str):
names = [n.strip() for n in names.split(",") if n.strip()]
return [build_algorithm(n, **kwargs) for n in names]
def algorithm_report() -> List[Dict[str, Any]]:
return [
{"name": k, "category": v["category"], "description": v["description"]}
for k, v in sorted(ALGORITHMS.items(), key=lambda kv: (kv[1]["category"], kv[0]))
]
def categories() -> List[str]:
return sorted({v["category"] for v in ALGORITHMS.values()})
__all__ = [
"ALGORITHMS",
"AlgorithmContext",
"BaseAlgorithm",
"EMA",
"SWA",
"Curriculum",
"DifficultyCurriculum",
"DPO",
"Distillation",
"LabelSmoothing",
"Mixup",
"WeightDecaySchedule",
"GradientAccumulation",
"GradientClip",
"GradientNoise",
"ZLoss",
"ActivationCheckpointing",
"SAM",
"Lookahead",
"algorithm_names",
"algorithm_report",
"build_algorithm",
"build_algorithms",
"categories",
"distillation_loss",
"dpo_loss",
"get_algorithm",
"sequence_log_probs",
]
|