experiments/shared/src/utils/parameter_count.py

"""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