torch-weighttracker 0.1.0__py3-none-any.whl

torch_weighttracker/__init__.py

@@ -0,0 +1,5 @@
+from torch_weighttracker.weight_tracker import WeightTracker
+
+__all__ = ["WeightTracker"]
+
+__version__ = "0.1.0"

torch_weighttracker/calculations/__init__.py

@@ -0,0 +1,77 @@
+from torch_weighttracker.calculations.base import (
+    BaseCalculation,
+    CalcType,
+    Calculation,
+)
+from torch_weighttracker.calculations.cached_calc import CachedCalculation
+from torch_weighttracker.calculations.calculations import (
+    CALCULATION_SPECS,
+    ActiveMacsPrModuleCalc,
+    BaselineMacsPrModuleCalc,
+    BaselineModuleAxesCalc,
+    BaselineParamPrUnitPrGroup,
+    CalculationContext,
+    CalculationSpec,
+    InitialUnitCountPrGroup,
+    ParamPrUnit,
+    UnitActiveMaskCalc,
+    UnitPrGroup,
+    create_active_macs_pr_module_calc,
+    create_active_units_calc,
+    create_baseline_group_sizes_calc,
+    create_baseline_macs_pr_module_calc,
+    create_baseline_module_axes_calc,
+    create_baseline_param_pr_unit_pr_group_calc,
+    create_bitrate_pr_module_calc,
+    create_group_change_effect_calc,
+    create_group_unit_param_change_calc,
+    create_group_sizes_calc,
+    create_groups_to_units_calc,
+    create_l2_norm_pr_unit_calc,
+    create_param_pr_unit_calc,
+    create_structured_unit_sum_calc,
+    create_unit_active_mask_calc,
+    create_unit_delta_to_module_axis_calc,
+    create_units_to_group_calc,
+    create_units_to_module_axis_calc,
+)
+from torch_weighttracker.calculations.pipeline_calc import PipelineCalc
+from torch_weighttracker.calculations.reduction_calc import ReductionCalc
+
+__all__ = [
+    "ActiveMacsPrModuleCalc",
+    "BaseCalculation",
+    "BaselineMacsPrModuleCalc",
+    "BaselineModuleAxesCalc",
+    "BaselineParamPrUnitPrGroup",
+    "CALCULATION_SPECS",
+    "CachedCalculation",
+    "CalcType",
+    "Calculation",
+    "CalculationContext",
+    "CalculationSpec",
+    "InitialUnitCountPrGroup",
+    "ParamPrUnit",
+    "PipelineCalc",
+    "ReductionCalc",
+    "UnitActiveMaskCalc",
+    "UnitPrGroup",
+    "create_active_macs_pr_module_calc",
+    "create_active_units_calc",
+    "create_baseline_group_sizes_calc",
+    "create_baseline_macs_pr_module_calc",
+    "create_baseline_module_axes_calc",
+    "create_baseline_param_pr_unit_pr_group_calc",
+    "create_bitrate_pr_module_calc",
+    "create_group_change_effect_calc",
+    "create_group_unit_param_change_calc",
+    "create_group_sizes_calc",
+    "create_groups_to_units_calc",
+    "create_l2_norm_pr_unit_calc",
+    "create_param_pr_unit_calc",
+    "create_structured_unit_sum_calc",
+    "create_unit_active_mask_calc",
+    "create_unit_delta_to_module_axis_calc",
+    "create_units_to_group_calc",
+    "create_units_to_module_axis_calc",
+]

torch_weighttracker/calculations/base.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from enum import Enum
+
+import torch
+from torch import nn
+
+
+class CalcType(str, Enum):
+    STRUCTURED_UNIT_SUM = "structured_unit_sum"
+    ACTIVE_UNITS = "active_units"
+    L2_NORM_PR_UNIT = "l2_norm_pr_unit"
+    BITRATE_PR_MODULE = "bitrate_pr_module"
+    UNITS_TO_MODULE_AXIS = "units_to_module_axis"
+    UNIT_DELTA_TO_MODULE_AXIS = "unit_delta_to_module_axis"
+    ACTIVE_MACS_PR_MODULE = "active_macs_pr_module"
+    BASELINE_MACS_PR_MODULE = "baseline_macs_pr_module"
+    BASELINE_MODULE_AXES = "baseline_module_axes"
+    UNITS_TO_GROUP = "units_to_group"
+    GROUPS_TO_UNITS = "groups_to_units"
+    UNIT_ACTIVE_MASK = "unit_active_mask"
+    GROUP_CHANGE_EFFECT = "group_change_effect"
+    GROUP_UNIT_PARAM_CHANGE = "group_unit_param_change"
+    BASELINE_PARAM_PR_UNIT_PR_GROUP = "baseline_param_pr_unit_pr_group"
+    PARAM_PR_UNIT = "param_pr_unit"
+    GROUP_SIZES = "group_sizes"
+    INIT_UNIT_PR_GROUP_COUNT = "baseline_group_sizes"
+
+
+class BaseCalculation(nn.Module):
+    calculation_type: CalcType | None = None
+
+    def __init__(
+        self,
+        dependencies: Mapping[CalcType, nn.Module] | None = None,
+    ) -> None:
+        super().__init__()
+        self.dependencies = nn.ModuleDict(
+            {}
+            if dependencies is None
+            else {calc_type.name: module for calc_type, module in dependencies.items()}
+        )
+        self.register_buffer("output_anchor", torch.empty(()), persistent=False)
+
+    def calc(self, calc_type: CalcType | str) -> nn.Module:
+        calc_type = CalcType(calc_type)
+        return self.dependencies[calc_type.name]
+
+    def compute(self, calc_type: CalcType | str, *args, **kwargs) -> torch.Tensor:
+        return self.calc(calc_type)(*args, **kwargs)
+
+
+Calculation = BaseCalculation

torch_weighttracker/calculations/cached_calc.py

@@ -0,0 +1,23 @@
+from typing import Any
+
+import torch
+from torch import nn
+
+
+class CachedCalculation(nn.Module):
+    def __init__(self, calculation: nn.Module) -> None:
+        super().__init__()
+        self.calculation = calculation
+        self.calculation_type = getattr(calculation, "calculation_type", None)
+
+        if not hasattr(calculation, "output_anchor"):
+            raise ValueError(
+                f"{calculation.__class__.__name__} must define output_anchor"
+            )
+
+    @torch.no_grad()
+    def refresh_cache(self, *args: Any, **kwargs: Any) -> None:
+        self.calculation.output_anchor = self.calculation(*args, **kwargs)
+
+    def forward(self) -> torch.Tensor:
+        return self.calculation.output_anchor

torch_weighttracker/calculations/calcs/__init__.py

@@ -0,0 +1,68 @@
+from torch_weighttracker.calculations.calcs.active_macs_pr_module import (
+    ActiveMacsPrModuleCalc,
+    create_active_macs_pr_module_calc,
+)
+from torch_weighttracker.calculations.calcs.active_units import create_active_units_calc
+from torch_weighttracker.calculations.calcs.baseline_macs_pr_module import (
+    BaselineMacsPrModuleCalc,
+    create_baseline_macs_pr_module_calc,
+)
+from torch_weighttracker.calculations.calcs.baseline_group_sizes import (
+    InitialUnitCountPrGroup,
+    create_baseline_group_sizes_calc,
+)
+from torch_weighttracker.calculations.calcs.baseline_module_axes import (
+    BaselineModuleAxesCalc,
+    create_baseline_module_axes_calc,
+)
+from torch_weighttracker.calculations.calcs.baseline_param_pr_unit_pr_group import (
+    BaselineParamPrUnitPrGroup,
+    create_baseline_param_pr_unit_pr_group_calc,
+)
+from torch_weighttracker.calculations.calcs.bitrate_pr_module import create_bitrate_pr_module_calc
+from torch_weighttracker.calculations.calcs.group_change_effect import create_group_change_effect_calc
+from torch_weighttracker.calculations.calcs.group_unit_param_change import create_group_unit_param_change_calc
+from torch_weighttracker.calculations.calcs.group_sizes import UnitPrGroup, create_group_sizes_calc
+from torch_weighttracker.calculations.calcs.groups_to_units import create_groups_to_units_calc
+from torch_weighttracker.calculations.calcs.l2_norm_pr_unit import (
+    L2NormPrUnit,
+    create_l2_norm_pr_unit_calc,
+)
+from torch_weighttracker.calculations.calcs.param_pr_unit import ParamPrUnit, create_param_pr_unit_calc
+from torch_weighttracker.calculations.calcs.structured_unit_sum import create_structured_unit_sum_calc
+from torch_weighttracker.calculations.calcs.unit_active_mask import UnitActiveMaskCalc, create_unit_active_mask_calc
+from torch_weighttracker.calculations.calcs.unit_delta_to_module_axis import (
+    create_unit_delta_to_module_axis_calc,
+)
+from torch_weighttracker.calculations.calcs.units_to_group import create_units_to_group_calc
+from torch_weighttracker.calculations.calcs.units_to_module_axis import create_units_to_module_axis_calc
+
+__all__ = [
+    "ActiveMacsPrModuleCalc",
+    "InitialUnitCountPrGroup",
+    "L2NormPrUnit",
+    "BaselineMacsPrModuleCalc",
+    "BaselineModuleAxesCalc",
+    "BaselineParamPrUnitPrGroup",
+    "ParamPrUnit",
+    "UnitPrGroup",
+    "UnitActiveMaskCalc",
+    "create_active_macs_pr_module_calc",
+    "create_active_units_calc",
+    "create_baseline_group_sizes_calc",
+    "create_baseline_macs_pr_module_calc",
+    "create_baseline_module_axes_calc",
+    "create_baseline_param_pr_unit_pr_group_calc",
+    "create_bitrate_pr_module_calc",
+    "create_group_change_effect_calc",
+    "create_group_unit_param_change_calc",
+    "create_group_sizes_calc",
+    "create_groups_to_units_calc",
+    "create_l2_norm_pr_unit_calc",
+    "create_param_pr_unit_calc",
+    "create_structured_unit_sum_calc",
+    "create_unit_active_mask_calc",
+    "create_unit_delta_to_module_axis_calc",
+    "create_units_to_group_calc",
+    "create_units_to_module_axis_calc",
+]

torch_weighttracker/calculations/calcs/active_macs_pr_module.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+import torch
+import torch.nn as nn
+
+from torch_weighttracker.calculations.base import CalcType, Calculation
+from torch_weighttracker.calculations.context import CalculationContext
+from torch_weighttracker.calculations.spec import CalculationSpec
+
+
+class ActiveMacsPrModuleCalc(Calculation):
+    """
+    Returns the active runtime MAC count for each weighted module.
+
+    Output: 1D tensor with length `len(weighted_modules)`.
+    Input: none.
+    """
+
+    calculation_type = CalcType.ACTIVE_MACS_PR_MODULE
+
+    def __init__(self, dependencies: Mapping[CalcType, nn.Module]) -> None:
+        super().__init__(dependencies)
+
+    def forward(self) -> torch.Tensor:
+        active_units = self.compute(CalcType.UNIT_ACTIVE_MASK)
+        baseline_axes = self.compute(CalcType.BASELINE_MODULE_AXES)
+        baseline_macs = self.compute(CalcType.BASELINE_MACS_PR_MODULE)
+        axis_delta = self.compute(
+            CalcType.UNIT_DELTA_TO_MODULE_AXIS,
+            active_units,
+        ).view_as(baseline_axes)
+        active_axes = baseline_axes + axis_delta
+        return baseline_macs * (active_axes / baseline_axes).prod(dim=1)
+
+
+def create_active_macs_pr_module_calc(
+    ctx: CalculationContext,
+    *,
+    dependencies: Mapping[CalcType, nn.Module],
+) -> ActiveMacsPrModuleCalc:
+    return ActiveMacsPrModuleCalc(dependencies)
+
+
+CALCULATION_SPEC = CalculationSpec(
+    calculation_type=CalcType.ACTIVE_MACS_PR_MODULE,
+    required_calculations=(
+        CalcType.UNIT_ACTIVE_MASK,
+        CalcType.UNIT_DELTA_TO_MODULE_AXIS,
+        CalcType.BASELINE_MACS_PR_MODULE,
+        CalcType.BASELINE_MODULE_AXES,
+    ),
+    create=lambda ctx, deps: create_active_macs_pr_module_calc(
+        ctx,
+        dependencies=deps,
+    ),
+)

torch_weighttracker/calculations/calcs/active_units.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+from torch_weighttracker.calculations.base import CalcType
+from torch_weighttracker.calculations.reduction_calc import ReductionCalc
+from torch_weighttracker.calculations.spec import CalculationSpec
+from torch_weighttracker.canonical_units import CanonicalUnitGroup
+from torch_weighttracker.operations import WeightOperationType
+from torch_weighttracker.plans.unit_weight_operation_plan import create_group_member_plan
+
+
+def create_active_units_calc(
+    groups: Iterable[CanonicalUnitGroup],
+) -> ReductionCalc:
+    """
+    Returns the active parameter count for each canonical unit.
+
+    Output: 1D tensor with length equal to the total canonical unit count.
+    Input: none.
+    """
+    return ReductionCalc(
+        create_group_member_plan(groups, WeightOperationType.ACTIVE),
+        calculation_type=CalcType.ACTIVE_UNITS,
+    )
+
+
+CALCULATION_SPEC = CalculationSpec(
+    calculation_type=CalcType.ACTIVE_UNITS,
+    create=lambda ctx, deps: create_active_units_calc(ctx.canonical_groups),
+)

torch_weighttracker/calculations/calcs/baseline_group_sizes.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import torch
+
+from torch_weighttracker.calculations.base import CalcType
+from torch_weighttracker.calculations.context import calculation_device, calculation_dtype
+from torch_weighttracker.calculations.spec import CalculationSpec
+from torch_weighttracker.calculations.static_calc import StaticCalc
+from torch_weighttracker.canonical_units import CanonicalUnitGroup
+
+
+class InitialUnitCountPrGroup(StaticCalc):
+    """
+    Returns the initial count of canonical units for each canonical group.
+
+    Output: 1D tensor with length `len(canonical_groups)`.
+    Input: none.
+    """
+
+    calculation_type = CalcType.INIT_UNIT_PR_GROUP_COUNT
+
+    def __init__(
+        self,
+        initial_unit_count_pr_group: torch.Tensor,
+    ) -> None:
+        super().__init__(initial_unit_count_pr_group)
+
+
+def create_baseline_group_sizes_calc(
+    groups: Iterable[CanonicalUnitGroup],
+    *,
+    device: torch.device | str,
+    dtype: torch.dtype,
+) -> InitialUnitCountPrGroup:
+    group_lengths = [group.length for group in groups]
+    return InitialUnitCountPrGroup(
+        torch.tensor(group_lengths, device=device, dtype=dtype)
+    )
+
+
+CALCULATION_SPEC = CalculationSpec(
+    calculation_type=CalcType.INIT_UNIT_PR_GROUP_COUNT,
+    cache_constant=True,
+    create=lambda ctx, deps: create_baseline_group_sizes_calc(
+        ctx.canonical_groups,
+        device=calculation_device(ctx),
+        dtype=calculation_dtype(ctx),
+    ),
+)

torch_weighttracker/calculations/calcs/baseline_macs_pr_module.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import torch
+
+from torch_weighttracker.calculations.base import CalcType
+from torch_weighttracker.calculations.context import (
+    CalculationContext,
+    calculation_device,
+    calculation_dtype,
+)
+from torch_weighttracker.calculations.spec import CalculationSpec
+from torch_weighttracker.calculations.static_calc import StaticCalc
+
+
+class BaselineMacsPrModuleCalc(StaticCalc):
+    """
+    Returns the initial runtime MAC count for each weighted module.
+
+    Output: 1D tensor with length `len(weighted_modules)`.
+    Input: none.
+    """
+
+    calculation_type = CalcType.BASELINE_MACS_PR_MODULE
+
+    def __init__(self, baseline_macs: torch.Tensor) -> None:
+        super().__init__(baseline_macs)
+
+
+def create_baseline_macs_pr_module_calc(
+    ctx: CalculationContext,
+    *,
+    device: torch.device | str,
+    dtype: torch.dtype,
+) -> BaselineMacsPrModuleCalc:
+    if ctx.example_inputs is None:
+        raise ValueError(
+            "BASELINE_MACS_PR_MODULE requires example_inputs so fvcore can "
+            "compute runtime MACs."
+        )
+
+    try:
+        from fvcore.nn import FlopCountAnalysis
+    except ImportError as error:
+        raise RuntimeError(
+            "BASELINE_MACS_PR_MODULE requires fvcore. Install fvcore or disable "
+            "structured BOPs MAC accounting."
+        ) from error
+
+    values = _fvcore_macs_by_weighted_module(ctx, FlopCountAnalysis)
+    return BaselineMacsPrModuleCalc(
+        torch.tensor(values, dtype=dtype, device=device)
+    )
+
+
+def _fvcore_macs_by_weighted_module(
+    ctx: CalculationContext,
+    flop_count_analysis,
+) -> list[float]:
+    names_by_module = {module: name for name, module in ctx.model.named_modules()}
+    analysis = flop_count_analysis(ctx.model, ctx.example_inputs)
+    if hasattr(analysis, "unsupported_ops_warnings"):
+        analysis = analysis.unsupported_ops_warnings(False)
+    if hasattr(analysis, "uncalled_modules_warnings"):
+        analysis = analysis.uncalled_modules_warnings(False)
+
+    by_module = analysis.by_module()
+    uncalled = (
+        set(analysis.uncalled_modules())
+        if hasattr(analysis, "uncalled_modules")
+        else set()
+    )
+
+    values: list[float] = []
+    missing: list[str] = []
+    for module in ctx.weighted_modules:
+        name = names_by_module.get(module)
+        if name is None or name not in by_module or name in uncalled:
+            missing.append("<unnamed>" if name is None else name)
+            continue
+        values.append(float(by_module[name]))
+
+    if missing:
+        raise ValueError(
+            "fvcore did not report MACs for weighted modules: "
+            + ", ".join(missing)
+        )
+
+    return values
+
+
+CALCULATION_SPEC = CalculationSpec(
+    calculation_type=CalcType.BASELINE_MACS_PR_MODULE,
+    requires_groups=False,
+    cache_constant=True,
+    create=lambda ctx, deps: create_baseline_macs_pr_module_calc(
+        ctx,
+        device=calculation_device(ctx),
+        dtype=calculation_dtype(ctx),
+    ),
+)

torch_weighttracker/calculations/calcs/baseline_module_axes.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+
+from torch_weighttracker.calculations.base import CalcType
+from torch_weighttracker.calculations.context import calculation_device, calculation_dtype
+from torch_weighttracker.calculations.spec import CalculationSpec
+from torch_weighttracker.calculations.static_calc import StaticCalc
+
+
+class BaselineModuleAxesCalc(StaticCalc):
+    """
+    Returns the initial input-axis and output-axis size for each weighted module.
+
+    Output: 2D tensor with shape `(len(weighted_modules), 2)`.
+    Input: none.
+    """
+
+    calculation_type = CalcType.BASELINE_MODULE_AXES
+
+    def __init__(self, baseline_axes: torch.Tensor) -> None:
+        super().__init__(baseline_axes)
+
+
+def create_baseline_module_axes_calc(
+    modules: Iterable[nn.Module],
+    *,
+    device: torch.device | str,
+    dtype: torch.dtype,
+) -> BaselineModuleAxesCalc:
+    baseline_axes = torch.tensor(
+        [_module_axis_sizes(module) for module in modules],
+        dtype=dtype,
+        device=device,
+    )
+    return BaselineModuleAxesCalc(baseline_axes)
+
+
+def _module_axis_sizes(module: nn.Module) -> tuple[float, float]:
+    if isinstance(module, nn.Conv2d):
+        return _conv_axis_sizes(module)
+
+    if isinstance(module, nn.Linear):
+        return float(module.in_features), float(module.out_features)
+
+    if isinstance(module, nn.MultiheadAttention):
+        raise ValueError(
+            "ACTIVE_MACS_PR_MODULE does not support nn.MultiheadAttention parent "
+            "modules in V1. Use projection Linear modules or add explicit MHA "
+            "operation terms."
+        )
+
+    if isinstance(module, nn.modules.batchnorm._BatchNorm):
+        return float(module.num_features), float(module.num_features)
+
+    if isinstance(module, nn.LayerNorm):
+        feature_dim = _layernorm_feature_dim(module)
+        return float(feature_dim), float(feature_dim)
+
+    if isinstance(module, nn.GroupNorm):
+        return float(module.num_channels), float(module.num_channels)
+
+    if isinstance(module, nn.modules.instancenorm._InstanceNorm):
+        return float(module.num_features), float(module.num_features)
+
+    if isinstance(module, nn.Embedding):
+        return float(module.embedding_dim), float(module.embedding_dim)
+
+    raise ValueError(
+        "Baseline module-axis sizes are not implemented for "
+        f"{module.__class__.__name__}."
+    )
+
+
+def _conv_axis_sizes(module: nn.Conv2d) -> tuple[float, float]:
+    if module.groups != 1:
+        raise ValueError(
+            "ACTIVE_MACS_PR_MODULE currently supports only Conv2d(groups=1)."
+        )
+    return float(module.in_channels), float(module.out_channels)
+
+
+def _layernorm_feature_dim(module: nn.LayerNorm) -> int:
+    normalized_shape = tuple(module.normalized_shape)
+    if len(normalized_shape) == 0:
+        raise ValueError("LayerNorm normalized_shape must not be empty.")
+    return int(normalized_shape[-1])
+
+
+CALCULATION_SPEC = CalculationSpec(
+    calculation_type=CalcType.BASELINE_MODULE_AXES,
+    requires_groups=False,
+    cache_constant=True,
+    create=lambda ctx, deps: create_baseline_module_axes_calc(
+        ctx.weighted_modules,
+        device=calculation_device(ctx),
+        dtype=calculation_dtype(ctx),
+    ),
+)