| """Parameter counting utilities.""" | |
| from __future__ import annotations | |
| import json | |
| from pathlib import Path | |
| from torch import nn | |
| def parameter_count_by_component(model: nn.Module) -> dict[str, int]: | |
| counts = {name: sum(p.numel() for p in module.parameters()) for name, module in model.named_children()} | |
| counts["total"] = sum(p.numel() for p in model.parameters()) | |
| return counts | |
| def save_parameter_count(model: nn.Module, path: str | Path) -> dict[str, int]: | |
| counts = parameter_count_by_component(model) | |
| out = Path(path) | |
| out.parent.mkdir(parents=True, exist_ok=True) | |
| out.write_text(json.dumps(counts, indent=2)) | |
| return counts | |