torch-weighttracker

Track, regularize, and prune structured units in PyTorch models.

torch-weighttracker gives you a model-level view of sparsity: channels, features, attention heads, head dimensions, and fused QKV slices are grouped into canonical units, so metrics, pruning, and regularizers operate on the structure you would actually compress.

import torch
import timm

from torch_weighttracker import WeightTracker
from torch_weighttracker.integrations.timm import infer_vit_num_heads

model = timm.create_model("vit_base_patch16_224", pretrained=False)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
example_inputs = torch.randn(1, 3, 224, 224)

tracker = WeightTracker(
    model,
    example_inputs=example_inputs,
    num_heads=infer_vit_num_heads(model),
    prune_num_heads=True,
)

tracker.create_tracker("structured_bops", log_total_bops=True)
group_lasso = tracker.create_regularizer("group_lasso")

for inputs, targets in dataloader:
    optimizer.zero_grad()

    outputs = model(inputs)
    task_loss = criterion(outputs, targets)
    loss = task_loss + 1e-4 * group_lasso()

    loss.backward()
    optimizer.step()

    metrics = tracker.track()

tracker.prune_zero_units()

Installation

python -m pip install torch-weighttracker

BOPs MAC accounting uses fvcore for baseline per-module MACs:

python -m pip install "torch-weighttracker[structured-bops]"

Why Use It?

PyTorch makes it easy to inspect one parameter tensor at a time. Structured compression often needs a different view:

• A channel can be coupled across convolutions, batch norms, linear layers, and residual paths.
• A transformer unit can mean an attention head, a head dimension, or a fused QKV slice rather than a simple row or column.
• A metric such as "active BOPs" depends on sparsity, module shape, MAC counts, and bit-widths at the same time.
• A regularizer such as group lasso should penalize the coupled structural unit, not each weight tensor independently.

WeightTracker turns those coupled structures into canonical units, then lets calculations operate over the canonical units with reusable tensor programs.

Use Cases

Current use cases:

• Add structured group lasso to a training loss.
• Track active structured or unstructured BOPs and compression rate during pruning, sparsity-aware training, or quantization-aware training (QAT).
• Inspect which modules participate in each channel, feature, head, or head-dim group.
• Build structural metrics that aggregate many weight tensors into one value per pruning unit.
• Physically prune zeroed canonical units, including attention heads, after sparsity-aware training.

Current Pruning Notes

WeightTracker can inspect zeroed canonical units with view_zero_units() / view_zero_structures() and physically remove them with prune_zero_units() / prune_zero_structures(). You can also remove one canonical unit directly with prune_unit(group_id, unit_id).

Zero detection can ignore module instances or module types with ignore_condition, matching tracker filter semantics. This only decides whether a structure is zero; if that structure is pruned, the coupled Torch-Pruning group is still applied. Use ignore_prune separately for modules or module types that should block the physical prune, such as a stem conv1:

zero_view = tracker.view_zero_structures(
    ignore_condition=[torch.nn.BatchNorm2d],
)
tracker.prune_zero_structures(
    ignore_condition=[torch.nn.BatchNorm2d],
    ignore_prune=[model.conv1],
)

Physical pruning changes module shapes and rebuilds the dependency state. Any registered trackers or regularizers are cleared after prune_unit() or prune_zero_units() / prune_zero_structures(), so recreate them before collecting metrics or losses from the pruned model:

metrics_before = tracker.create_tracker(
    "structured_bops",
    log_total_bops=True,
).track()

tracker.prune_zero_units()

metrics_after = tracker.create_tracker(
    "structured_bops",
    log_total_bops=True,
).track()

Fake pruning remains useful during training because it zeros the selected canonical unit while keeping the module shapes intact:

tracker.create_tracker("structured_bops", log_total_bops=True)
metrics_before = tracker.track()

tracker.fake_prune_unit(group_id=3, unit_id=0)
tracker.fake_prune_unit(group_id=3, unit_id=2)

metrics_after = tracker.track()

timm ViT Attention Heads

The torch_weighttracker.integrations.timm helpers make timm ViT attention blocks visible as head-level pruning groups. infer_vit_num_heads(model) maps each fused Attention.qkv projection to its current head count, and sync_vit_attention_metadata updates timm attention metadata after physical head pruning.

import timm
import torch

from torch_weighttracker import WeightTracker
from torch_weighttracker.integrations.timm import (
    infer_vit_num_heads,
    sync_vit_attention_metadata,
)

example_inputs = torch.rand(1, 3, 224, 224)
model = timm.create_model(
    "vit_base_patch16_224",
    pretrained=False,
    num_classes=10,
)

tracker = WeightTracker(
    model,
    example_inputs=example_inputs,
    num_heads=infer_vit_num_heads(model),
    prune_num_heads=True,
    post_prune_hooks=(sync_vit_attention_metadata,),
)

print(tracker.view_structures())

tracker.create_tracker("structured_bops", log_total_bops=True)
metrics_before = tracker.track()

# Example: zero two attention heads in the group reported by view_structures().
tracker.fake_prune_unit(group_id=3, unit_id=0)
tracker.fake_prune_unit(group_id=3, unit_id=2)
metrics_after_fake_prune = tracker.track()

# Convert zeroed units into real shape changes, then recreate metric trackers.
tracker.prune_zero_units()
metrics_after_physical_prune = tracker.create_tracker(
    "structured_bops",
    log_total_bops=True,
).track()

print(metrics_before["structured_bops"])
print(metrics_after_fake_prune["structured_bops"])
print(metrics_after_physical_prune["structured_bops"])

For timm ViTs, head pruning removes complete q/k/v head slices from the fused qkv projection and the corresponding projection input channels. The sync hook keeps num_heads, attn_dim, head_dim, and scale consistent with the new shape so the pruned model can still run a forward pass.

Hugging Face BERT Heads and FFN Neurons

The optional BERT adapter keeps encoder self-attention heads and intermediate FFN neurons as the only canonical units. Shared hidden-width groups and the QA head are excluded from tracking, regularization, zero detection, and pruning. Transformers remains a development-only dependency; applications can install a compatible v4 release separately.

from torch_weighttracker import WeightTracker
from torch_weighttracker.integrations.transformers import bert_pruning_config

tracker = WeightTracker(
    model,
    example_inputs=(input_ids, attention_mask, token_type_ids),
    **bert_pruning_config(model),
)

# Head groups use Hugging Face's model-level callback. FFN groups continue to
# use Torch-Pruning. Both paths rebuild tracker state after physical pruning.
tracker.fake_prune_unit(group_id=0, unit_id=1)
tracker.prune_zero_structures()

SeparateQKVAttentionSpec is the framework-neutral API behind the adapter. Its callback receives positions in the currently compact layer; an integration is responsible for mapping those positions to any framework-specific stable IDs. ignore_prune protects both the owning attention module and every constituent Q/K/V/output projection.

The BERT adapter requires the model-level prune_heads API available in Transformers v4. Transformers v5 removed head pruning, so the adapter fails with a clear compatibility error instead of providing a version-specific compactor.

group_pruning_summary preserves the weight-only meaning of pruned_params and adds pruned_bias_params plus their sum, pruned_physical_params, both globally and per group. Biases do not affect PARAM_PR_UNIT, L2, or GroupLasso.

The deterministic proof notebook is sanity_checks/bert_pruning_playground.ipynb. It reports group membership, L2 and active-unit heatmaps, shape and physical-parameter changes, weighted-layer MAC/BOP changes, unsupported-operation warnings, and fake-versus-compact logit parity. Its BOP totals are explicitly weighted-layer estimates when attention score operations fall outside fvcore coverage.

Group Lasso

Structured group lasso regularizes coupled units together. Layers can be excluded per regularizer:

from torch_weighttracker.regularizers import RegularizerType

group_lasso = tracker.create_regularizer(
    RegularizerType.GROUP_LASSO,
    ignore=[model.classifier],
)

loss = task_loss + 1e-4 * group_lasso()
loss.backward()

Structured BOPs

Structured BOPs reports compression against a dense 32-bit baseline by default. BOP trackers exclude normalization modules such as batchnorm, layernorm, groupnorm, and instancenorm layers by default:

import torch

from torch_weighttracker.trackers import TrackerType

metrics = tracker.create_tracker(
    TrackerType.STRUCTURED_BOPS,
    include=[model.layer3, model.layer4],
).track()

print(metrics["structured_bops"]["compression"])

raw_metrics = tracker.create_tracker(
    TrackerType.STRUCTURED_BOPS,
    include=[model.layer3, model.layer4],
    log_total_bops=True,
    log_layerwise_stats=True,
).track()

structured = raw_metrics["structured_bops"]
print(structured["bops"])
print(structured["modules"])

If you want to report compression against an earlier model, pass that per-module dense MAC vector when creating the BOP tracker. Values are raw MACs, ordered after this tracker's include/ignore filters. This only changes metric normalization and logged baselines; active BOPs still come from the current model:

structured_bops = tracker.create_tracker(
    TrackerType.STRUCTURED_BOPS,
    ignore=ignored_layers,
    normalization_macs_pr_module=original_dense_macs,
)

create_tracker accepts a single TrackerType/string or a list of tracker types/strings:

tracker.create_tracker(
    [TrackerType.STRUCTURED_BOPS, "group_pruning_summary"]
)
metrics = tracker.track()

Tracker metrics are nested by tracker name and convert tensor values to Python numbers/lists by default. Pass convert_tensors=False to preserve tensors in the returned metrics. Pass wandb_format=True to return a flat W&B-style dict:

tracker.create_tracker("structured_bops", wandb_format=True)
metrics = tracker.track()
# {"structured_bops/compression": 0.98}
# wandb.log(metrics)

Formulation of the Structured BOPs Metric

For each weighted module $m$, WeightTracker multiplies the active structured MAC count by the bit widths of that module's input activations and weights [1]:

$$ \mathit{StructuredBOPs}_m = \mathit{ActiveMACs}_m \cdot b^{\mathrm{act}}_m \cdot b^{\mathrm{weight}}_m $$

The active MAC count scales the dense module MAC count by the active fraction of each structural cost axis:

$$ \mathit{ActiveMACs}m = \mathit{BaselineMACs}m \cdot \prod{a \in A_m} \frac{n^{\mathrm{active}}{m,a}}{n^{\mathrm{baseline}}_{m,a}} $$

Compression is reported against a dense 32-bit activation and 32-bit weight baseline:

$$ \mathit{BaselineBOPs}_m = \mathit{BaselineMACs}_m \cdot 32 \cdot 32 $$

$$ \mathit{CompressionRate} = 1 - \frac{\sum_m \mathit{StructuredBOPs}_m} {\sum_m \mathit{BaselineBOPs}_m} $$

Where:

• $\mathit{StructuredBOPs}_m$: active bit operations for weighted module $m$.
• $\mathit{ActiveMACs}_m$: active MAC count after structured units are masked or pruned.
• $\mathit{BaselineMACs}_m$: dense MAC count for module $m$ before structured pruning.
• $A_m$: structural cost axes for module $m$, such as input and output channel axes.
• $n^{\mathrm{active}}_{m,a}$: active size of cost axis $a$ for module $m$.
• $n^{\mathrm{baseline}}_{m,a}$: dense baseline size of cost axis $a$ for module $m$.
• $b^{\mathrm{act}}_m$: bit width of the input activations of module $m$.
• $b^{\mathrm{weight}}_m$: weight bit width for module $m$.

Bit widths come from weight_bitwidth_fn and activation_bitwidth_fn; see Exposing Bit-Widths.

Comparison with Direct Removal and FLOP Count

For some model architectures, the BOPs calculation may differ from values reported by other libraries. These differences mainly come from which layers and operations are included. WeightTracker does not count elementwise operations such as ReLU activations or bias terms.

The repository includes sanity notebooks comparing fvcore.FlopCountAnalysis on physically pruned models with WeightTracker on fake-pruned models, where weights are zeroed to match the equivalent hard-pruned structure.

Local sanity notebooks compare WeightTracker MAC accounting with physically pruned models from Torch-Pruning. These dependencies are optional and are not installed with the base package:

python -m pip install -e ".[dev-local]"

Then start Jupyter from the repository root and open the notebooks in sanity_checks/.

Unstructured Sparsity

Unstructured sparsity reports exact zero-weight fractions. The total is weighted by each layer's number of weight elements, not averaged across layer fractions:

import torch

from torch_weighttracker.trackers import TrackerType

metrics = tracker.create_tracker(
    TrackerType.UNSTRUCTURED_SPARSITY,
    include=[model.layer3, model.layer4],
    ignore=[torch.nn.BatchNorm2d],
).track()

sparsity = metrics["unstructured_sparsity"]
print(sparsity["sparsity"])
print(sparsity["layers"])

Values are fractions in [0, 1]. Parametrized fake quantization is measured through the effective module.weight, so quantized zeros count as sparse weights.

Unstructured BOPs

Unstructured BOPs combines each layer's dense runtime MAC count with its active weight fraction and input-activation/weight bit widths:

import torch

from torch_weighttracker.trackers import TrackerType

metrics = tracker.create_tracker(
    TrackerType.UNSTRUCTURED_BOPS,
    include=[model.layer3, model.layer4],
).track()

print(metrics["unstructured_bops"]["compression"])

raw_metrics = tracker.create_tracker(
    TrackerType.UNSTRUCTURED_BOPS,
    include=[model.layer3, model.layer4],
    log_total_bops=True,
    log_layerwise_stats=True,
).track()

unstructured = raw_metrics["unstructured_bops"]
print(unstructured["bops"])
print(unstructured["modules"])

normalization_macs_pr_module is also supported here and follows the same filtered weighted-module order as STRUCTURED_BOPS.

For each weighted module $m$:

$$ \mathit{UnstructuredBOPs}_m = \mathit{BaselineMACs}_m \cdot (1 - \mathit{Sparsity}_m) \cdot b^{\mathrm{act}}_m \cdot b^{\mathrm{weight}}_m $$

Bit widths are set as for structured BOPs; see Exposing Bit-Widths. Compression uses the same dense 32-bit baseline as structured BOPs:

$$ \mathit{CompressionRate} = 1 - \frac{\sum_m \mathit{UnstructuredBOPs}_m} {\sum_m \mathit{BaselineMACs}_m \cdot 32 \cdot 32} $$

Exposing Bit-Widths

Both BOPs trackers multiply each layer's MACs by the bit-widths of the two operands in each multiplication: the layer's weights $b^{\mathrm{weight}}_m$ and the layer's input activations $b^{\mathrm{act}}_m$. This is the usual BOPs definition in the literature:

• Baskin et al. [3] count one $b_a$-bit by $b_w$-bit multiplication per MAC, plus accumulator additions.
• van Baalen et al. [4] use the same count without the accumulator terms, since the accumulator width is usually fixed.
• Wang et al. [1] write $\mathit{BOPs}l = \mathit{MACs}l \cdot b{w,l} \cdot b{a,l-1}$. The activation bit-width there is that of the previous layer's output, which is layer $l$'s input.

Each bit-width comes from its own function, passed when constructing WeightTracker:

def weight_bitwidth_fn(name: str, module: nn.Module) -> int:
    ...  # bit-width of this layer's weights


def activation_bitwidth_fn(name: str, module: nn.Module) -> int:
    ...  # bit-width of this layer's input activations


tracker = WeightTracker(
    model,
    example_inputs,
    weight_bitwidth_fn=weight_bitwidth_fn,
    activation_bitwidth_fn=activation_bitwidth_fn,
)

• name is the layer's qualified name from model.named_modules(), the same key used under modules in layerwise BOPs metrics. module is the layer.
• Both functions are called for every layer the BOPs trackers count, on every track(). Bit-width changes made by a quantizer during training show up in the next track(); trackers do not need to be recreated.
• Either function can be left out, and that side is then 32 bits. For weight-only quantization, pass only weight_bitwidth_fn. Without either function, BOPs only reflect sparsity.
• Values must be whole numbers >= 1: an int, a whole-number float, or a scalar tensor such as torch.tensor(4.). Anything else, such as 5.3, 0, or a boolean, raises a ValueError naming the layer. Round learned bit-widths in the quantizer or in the function.
• Bit-widths come only from these functions. The tracker does not read module attributes or quantizer parametrizations on its own.

Input Activations

activation_bitwidth_fn is called with the layer that consumes the activations, so an activation quantizer counts for the layer it feeds. In conv1 -> ReLU -> act_quant -> conv2, the bit-width of act_quant is the activation bit-width of conv2, not conv1.

The tracker cannot make this link itself, because quantizer modules are not visible in its dependency graph. Map it in the function instead, for example with a dict from layer name to the quantizer in front of that layer. For the first layer, return the bit-width of the network input, for example 8 for 8-bit images.

nn.MultiheadAttention counts as one layer with one activation and one weight bit-width. Its MACs cover only the input and output projections. The matrix products between activations ($QK^\top$ and attention times $V$) are not counted.

Example

weight_bits = {"conv1": 4, "conv2": 8}
input_bits = {"conv2": 8}

tracker = WeightTracker(
    model,
    example_inputs,
    weight_bitwidth_fn=lambda name, module: weight_bits.get(name, 32),
    activation_bitwidth_fn=lambda name, module: input_bits.get(name, 32),
)
bops = tracker.create_tracker("structured_bops")
bops.track()

weight_bits["conv1"] = 2  # e.g. a quantizer lowered conv1 during training
bops.track()  # uses the new bit-widths

In practice the functions read the current values from your quantizers instead of from dicts.

Inspecting Bit-Widths

view_bitwidths() shows what the trackers see. It accepts the same include/ignore filters as the BOPs trackers. Each source is the function name, or "default" when that function was not passed. For the example above:

for layer in tracker.view_bitwidths().layers:
    print(
        layer.name,
        layer.activation_bits,
        layer.weight_bits,
        layer.activation_source,
        layer.weight_source,
    )
# conv1 32 2 activation_bitwidth_fn weight_bitwidth_fn
# conv2 8 8 activation_bitwidth_fn weight_bitwidth_fn

Cost

Each function runs once per layer on every track(). If a function reads bit-widths from GPU tensors, each read is a GPU-to-CPU sync, so call track() at logging frequency rather than on every optimizer step.

python -m sanity_checks.bitwidth_recheck shows BOPs following a weight and an input-activation bit-width change between two track() calls.

Migrating from 0.3

Earlier versions read bit-widths once, when the tracker was created, from quantizer parametrizations on weight exposing get_bitwidth() and from the activation_bitrate, act_bitrate, weight_bitrate, and bitrate module attributes. These are no longer read. Pass the functions instead, for example:

tracker = WeightTracker(
    model,
    example_inputs,
    activation_bitwidth_fn=lambda name, m: getattr(m, "activation_bitrate", 32),
    weight_bitwidth_fn=lambda name, m: getattr(m, "weight_bitrate", 32),
)

NVIDIA 2:4 Sparsity

NVIDIA 2:4 sparsity reports block eligibility for supported weighted layers. Linear and MultiheadAttention projection weights are grouped in contiguous blocks of four along the input axis. Convolution weights shaped [K, C, ...] are grouped along C for each output/spatial position.

import torch

from torch_weighttracker.trackers import TrackerType

metrics = tracker.create_tracker(
    TrackerType.NVIDIA_2_4_SPARSITY,
    include=[model.layer3, model.layer4],
    ignore=[torch.nn.BatchNorm2d],
    log_layerwise_stats=True,
).track()

nvidia_24 = metrics["nvidia_2_4_sparsity"]
print(nvidia_24["strict_block_fraction"])
print(nvidia_24["nvidia_eligible_block_fraction"])
print(nvidia_24["tail_elements"])

The strict fraction counts complete 4-value blocks with exactly two zeros. The NVIDIA-eligible fraction counts blocks with at least two zeros, matching the TensorRT eligibility rule. Tail elements are reported separately and prevent a layer from counting as strict or eligible.

Group Pruning Summary

Group pruning summary reports pruned canonical units and group-attributed pruned parameters in a nested tracker metrics dictionary:

import torch

from torch_weighttracker.trackers import TrackerType

metrics = tracker.create_tracker(
    TrackerType.GROUP_PRUNING_SUMMARY,
    include=[model.layer3, model.layer4],
    ignore=[torch.nn.BatchNorm2d],
).track()

summary = metrics["group_pruning_summary"]
print(summary["pruned_units"])
print(summary["pruned_params"])

Per-group values are emitted under summary["groups"]["layer3.0.conv1:prune_out_channels"]["pruned_units"] and the corresponding "pruned_params" key.

Architecture

The main API is WeightTracker. Internally it is split into a few layers:

1. Dependency discovery: WeightTracker builds dependency groups from the model and example_inputs using Torch-Pruning's dependency graph machinery [2], whose work we gratefully build on.
2. Canonical units: canonical_units.py normalizes raw dependency groups into CanonicalUnitGroup objects. These give channels, features, attention heads, and head dimensions a shared unit index.
3. Reduction plans: reductions/ and plans/ compile module and unit mappings into segment and index operations that use PyTorch's efficient tensor computations.
4. Calculations: calculations/ defines named calculation specs such as per-unit L2 norm, active units, parameters per unit, active MACs, and bit-widths. Calculations can depend on each other and cache constant results.
5. Consumers: regularizers/ and trackers/ request the calculations they need, optionally with include and ignore contexts for selecting modules in a specific metric or regularizer.

The result is a small public surface with a reusable internal graph:

model + example inputs
        |
        v
dependency groups -> canonical units -> reduction plans -> calculations
                                                          |
                                                          v
                                               regularizers and trackers

Speed

Compared with a naive implementation, the current implementation gives the following speedups on ResNet 20 on a RTX 3060:

Sparse reduction backend

Reduction plans run on a sparse backend by default. Supported ops (sum, squared sum, mean, count, active, L1, L2, and affine pipeline steps) are lowered into a few sparse matrix-vector products over one flattened source vector, instead of one reduction plus one scatter per module. Ops it cannot express, such as fused QKV reductions, and source tensors larger than 2^18 elements stay on the reference loop path inside the same calculation, so results are unchanged.

Speedup over the loop backend on an RTX 3060 (sanity_checks/sparse_reduction_benchmark.py):

Select the backend per tracker or for a scope:

from torch_weighttracker import WeightTracker, use_reduction_backend

tracker = WeightTracker(model, example_inputs=x, reduction_backend="loop")

with use_reduction_backend("loop"):
    ...  # calculations constructed here use the loop backend

sanity_checks/sparse_reductions_playground.ipynb walks through the lowering on hand-checkable examples.

Status

This package is pre-1.0. Public APIs may still change while the tracker, calculation, and regularizer surfaces settle.

Future Work

1. Streamline definitions and method names across the codebase.
2. Improve calculation caching so repeated computations are not performed twice.
3. Improve compilation of computation plans for bigger speedups.
4. Improve memory management within calculations.
5. Write more comprehensive docstrings.

Future custom use cases will need a broader top-level WeightTracker API for custom operations, custom layers, and generic group definitions.

License

MIT

References

[1] Wang et al., Differentiable Joint Pruning and Quantization for Hardware Efficiency, ECCV 2020.

[2] Fang et al., Torch-Pruning.

[3] Baskin et al., UNIQ: Uniform Noise Injection for Non-Uniform Quantization of Neural Networks, 2018.

[4] van Baalen et al., Bayesian Bits: Unifying Quantization and Pruning, NeurIPS 2020.