File size: 7,519 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 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | """Shared helpers for the Agron neural-network zoo.
Everything here is deliberately dependency-light (torch + numpy only) so the
package can be imported from the Gradio app, the training scripts and the
test-suite without pulling in the full Hugging Face stack.
"""
from __future__ import annotations
import json
import math
import os
import random
from typing import Any, Dict, Iterable, List, Optional
import numpy as np
import torch
import torch.nn as nn
# --------------------------------------------------------------------------- #
# Reproducibility
# --------------------------------------------------------------------------- #
def set_seed(seed: int = 42, deterministic: bool = True) -> None:
"""Seed python / numpy / torch RNGs."""
random.seed(seed)
np.random.seed(seed % (2**32))
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
if deterministic:
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# --------------------------------------------------------------------------- #
# Model introspection
# --------------------------------------------------------------------------- #
def count_parameters(model: nn.Module, trainable_only: bool = False) -> int:
"""Number of (trainable) parameters in ``model``."""
return sum(p.numel() for p in model.parameters() if p.requires_grad or not trainable_only)
def format_parameters(n: int) -> str:
"""Human readable parameter count, e.g. ``50.3M``."""
if n >= 1e9:
return f"{n / 1e9:.2f}B"
if n >= 1e6:
return f"{n / 1e6:.1f}M"
if n >= 1e3:
return f"{n / 1e3:.1f}K"
return str(n)
def model_device(model: nn.Module) -> torch.device:
"""Best guess of the device a module lives on."""
try:
return next(model.parameters()).device
except StopIteration:
return torch.device("cpu")
def dtype_of(model: nn.Module) -> torch.dtype:
try:
return next(model.parameters()).dtype
except StopIteration:
return torch.float32
# --------------------------------------------------------------------------- #
# Initialisation
# --------------------------------------------------------------------------- #
def init_weights(module: nn.Module, std: float = 0.02) -> None:
"""Deterministic weight init used by every Agron architecture.
Linear layers get a truncated-normal style normal init, embeddings get a
normal init, and any residual-projection weight is zero-initialised so the
network starts as an identity map on the residual stream (a well known
trick that stabilises deep transformers).
"""
if isinstance(module, nn.Linear):
module.weight.data.normal_(mean=0.0, std=std)
if module.bias is not None:
module.bias.data.zero_()
elif isinstance(module, nn.Embedding):
module.weight.data.normal_(mean=0.0, std=std)
if module.padding_idx is not None:
module.weight.data[module.padding_idx].zero_()
elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
if module.weight is not None:
module.weight.data.fill_(1.0)
if module.bias is not None:
module.bias.data.zero_()
# --------------------------------------------------------------------------- #
# Small maths helpers
# --------------------------------------------------------------------------- #
def cosine_schedule(step: int, total: int, warmup: int = 0, min_ratio: float = 0.0) -> float:
"""Linear warmup followed by cosine decay to ``min_ratio``."""
if warmup > 0 and step < warmup:
return (step + 1) / max(1, warmup)
progress = (step - warmup) / max(1, total - warmup)
progress = min(max(progress, 0.0), 1.0)
return min_ratio + (1.0 - min_ratio) * 0.5 * (1.0 + math.cos(math.pi * progress))
def linear_schedule(step: int, total: int, warmup: int = 0, min_ratio: float = 0.0) -> float:
"""Linear warmup followed by linear decay to ``min_ratio``."""
if warmup > 0 and step < warmup:
return (step + 1) / max(1, warmup)
progress = (step - warmup) / max(1, total - warmup)
progress = min(max(progress, 0.0), 1.0)
return max(min_ratio, 1.0 - progress * (1.0 - min_ratio))
def warmup_constant(step: int, total: int, warmup: int = 0, min_ratio: float = 0.0) -> float:
"""Warmup then hold at 1.0 (times ``min_ratio`` afterwards)."""
if warmup > 0 and step < warmup:
return (step + 1) / max(1, warmup)
return 1.0
def inverse_sqrt_schedule(step: int, total: int, warmup: int = 0, min_ratio: float = 0.0) -> float:
"""Noam / Transformer-XL style ``1/sqrt(step)`` decay."""
step = max(step, 1)
if warmup > 0 and step < warmup:
return (step + 1) / max(1, warmup)
return min(1.0, math.sqrt(warmup / step) if warmup else 1.0)
def polynomial_schedule(step: int, total: int, warmup: int = 0, power: float = 1.0,
min_ratio: float = 0.0) -> float:
"""Polynomial (e.g. ``power=0.5``) decay."""
if warmup > 0 and step < warmup:
return (step + 1) / max(1, warmup)
progress = (step - warmup) / max(1, total - warmup)
progress = min(max(progress, 0.0), 1.0)
return max(min_ratio, (1.0 - progress) ** power)
def cosine_restart_schedule(step: int, total: int, warmup: int = 0, cycles: int = 1,
min_ratio: float = 0.0) -> float:
"""Cosine schedule with ``cycles`` warm restarts (SGDR)."""
if warmup > 0 and step < warmup:
return (step + 1) / max(1, warmup)
cycles = max(1, cycles)
cycle_len = max(1, (total - warmup) // cycles)
pos = (step - warmup) % cycle_len
progress = pos / cycle_len
return min_ratio + (1.0 - min_ratio) * 0.5 * (1.0 + math.cos(math.pi * progress))
def wsd_schedule(step: int, total: int, warmup: int = 0, min_ratio: float = 0.0,
decay_ratio: float = 0.1) -> float:
"""Warmup-Stable-Decay schedule."""
if warmup > 0 and step < warmup:
return (step + 1) / max(1, warmup)
stable_end = int(total * 0.8)
if step <= stable_end:
return 1.0
progress = (step - stable_end) / max(1, total - stable_end)
progress = min(max(progress, 0.0), 1.0)
return max(min_ratio, decay_ratio + (1.0 - decay_ratio) * (1.0 - progress))
# --------------------------------------------------------------------------- #
# JSON / IO
# --------------------------------------------------------------------------- #
def write_json(path: str, payload: Any) -> None:
"""Atomically write ``payload`` as JSON to ``path``."""
os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
tmp = f"{path}.tmp"
with open(tmp, "w", encoding="utf-8") as fh:
json.dump(payload, fh, indent=2, default=str)
os.replace(tmp, path)
def read_json(path: str, default: Any = None) -> Any:
"""Read JSON from ``path`` returning ``default`` on any failure."""
try:
with open(path, "r", encoding="utf-8") as fh:
return json.load(fh)
except Exception:
return default
def truncate_to_blocks(token_ids: List[int], block_size: int) -> List[List[int]]:
"""Split a flat token stream into fixed-size blocks."""
if block_size <= 0:
raise ValueError("block_size must be positive")
blocks = []
for i in range(0, max(0, len(token_ids) - block_size + 1), block_size):
blocks.append(token_ids[i:i + block_size])
return blocks
|