| from __future__ import annotations |
|
|
| import json |
| import math |
| import random |
| from copy import deepcopy |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from torch import Tensor, nn |
|
|
|
|
| def seed_everything(seed: int) -> None: |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
|
|
|
|
| def move_targets(targets: list[dict[str, Tensor]], device: torch.device): |
| return [ |
| {key: value.to(device, non_blocking=True) for key, value in target.items()} |
| for target in targets |
| ] |
|
|
|
|
| class ModelEMA: |
| def __init__(self, model: nn.Module, decay: float = 0.9998) -> None: |
| self.model = deepcopy(model).eval() |
| self.decay = decay |
| for parameter in self.model.parameters(): |
| parameter.requires_grad_(False) |
| target = dict(self.model.state_dict()) |
| self._float_names = [name for name, value in target.items() if value.is_floating_point()] |
| self._float_targets = [target[name] for name in self._float_names] |
| self._other = [(name, target[name]) for name in target if not target[name].is_floating_point()] |
| self._float_sources: list[Tensor] | None = None |
| self._other_sources: list[Tensor] | None = None |
|
|
| @torch.no_grad() |
| def update(self, model: nn.Module) -> None: |
| if self._float_sources is None: |
| source = dict(model.state_dict()) |
| self._float_sources = [source[name].detach() for name in self._float_names] |
| self._other_sources = [source[name].detach() for name, _ in self._other] |
| torch._foreach_mul_(self._float_targets, self.decay) |
| torch._foreach_add_(self._float_targets, self._float_sources, alpha=1.0 - self.decay) |
| for (_, target_value), source_value in zip(self._other, self._other_sources): |
| target_value.copy_(source_value) |
|
|
|
|
| def learning_rate_factor(step: int, total_steps: int, warmup_steps: int, min_ratio: float) -> float: |
| if warmup_steps > 0 and step < warmup_steps: |
| return max(step + 1, 1) / warmup_steps |
| progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1) |
| cosine = 0.5 * (1.0 + math.cos(math.pi * min(max(progress, 0.0), 1.0))) |
| return min_ratio + (1.0 - min_ratio) * cosine |
|
|
|
|
| def save_checkpoint(path: str | Path, **state) -> None: |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_suffix(path.suffix + ".tmp") |
| torch.save(state, temporary) |
| temporary.replace(path) |
|
|
|
|
| def write_json(path: str | Path, value) -> None: |
| with Path(path).open("w", encoding="utf-8") as handle: |
| json.dump(value, handle, indent=2) |
|
|
|
|
| def trainable_parameter_count(model: nn.Module) -> int: |
| return sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad) |
|
|
