batchnorm

batchnorm

BatchNorm detection and replacement for quantization preparation.

has_batchnorm_buffers

has_batchnorm_buffers(module)

Check if module has the four canonical BN attributes.

Checks for running_mean, running_var, weight, and bias as either buffers or parameters, since standard BN stores weight/bias as parameters while frozen variants store them as buffers. All four must be one-dimensional tensors of the same length, which rules out modules that merely reuse these attribute names.

Source code in src/versatil/post_training_compression/preparation/batchnorm.py

def has_batchnorm_buffers(module: nn.Module) -> bool:
    """Check if module has the four canonical BN attributes.

    Checks for running_mean, running_var, weight, and bias as either
    buffers or parameters, since standard BN stores weight/bias as
    parameters while frozen variants store them as buffers. All four must
    be one-dimensional tensors of the same length, which rules out modules
    that merely reuse these attribute names.
    """
    tensors: list[torch.Tensor] = []
    for name in BATCHNORM_ATTRIBUTE_NAMES:
        attribute = getattr(module, name, None)
        if not isinstance(attribute, torch.Tensor):
            return False
        tensors.append(attribute)
    return all(
        tensor.dim() == 1 and tensor.shape == tensors[0].shape for tensor in tensors
    )

is_frozen_batchnorm

is_frozen_batchnorm(module)

Check if module is a frozen or non-standard BN that needs replacement.

Source code in src/versatil/post_training_compression/preparation/batchnorm.py

def is_frozen_batchnorm(module: nn.Module) -> bool:
    """Check if module is a frozen or non-standard BN that needs replacement."""
    if not has_batchnorm_buffers(module):
        return False
    if _is_standard_batchnorm_already_prepared(module):
        return False
    return not isinstance(module, STANDARD_BATCHNORM_TYPES)

extract_batchnorm_parameters

extract_batchnorm_parameters(batchnorm)

Extract (running_mean, running_var, weight, bias, eps) from any BN variant.

Source code in src/versatil/post_training_compression/preparation/batchnorm.py

def extract_batchnorm_parameters(
    batchnorm: nn.Module,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, float] | None:
    """Extract (running_mean, running_var, weight, bias, eps) from any BN variant."""
    if not has_batchnorm_buffers(batchnorm):
        return None
    running_mean = batchnorm.running_mean
    running_var = batchnorm.running_var
    weight = batchnorm.weight
    bias = batchnorm.bias
    eps = getattr(batchnorm, "eps", 1e-5)
    return running_mean, running_var, weight, bias, eps

extract_activation

extract_activation(batchnorm)

Extract fused activation from BN modules like FrozenBatchNormAct2d.

Source code in src/versatil/post_training_compression/preparation/batchnorm.py

def extract_activation(batchnorm: nn.Module) -> nn.Module | None:
    """Extract fused activation from BN modules like FrozenBatchNormAct2d."""
    activation = getattr(batchnorm, "act", None)
    if activation is None:
        return None
    if isinstance(activation, nn.Identity):
        return None
    return activation

replace_frozen_batchnorm

replace_frozen_batchnorm(model)

Recursively replace all frozen BN with standard BatchNorm variants.

Returns the number of modules replaced.

Source code in src/versatil/post_training_compression/preparation/batchnorm.py

def replace_frozen_batchnorm(model: nn.Module) -> int:
    """Recursively replace all frozen BN with standard BatchNorm variants.

    Returns the number of modules replaced.
    """
    replacement_count = 0
    for name, child in list(model.named_children()):
        if is_frozen_batchnorm(child):
            num_features = child.running_mean.shape[0]
            replacement = _create_replacement_batchnorm(
                batchnorm=child,
                num_features=num_features,
            )
            activation = extract_activation(child)
            if activation is not None:
                combined = nn.Sequential(replacement, activation)
                setattr(model, name, combined)
            else:
                setattr(model, name, replacement)
            replacement_count += 1
        else:
            replacement_count += replace_frozen_batchnorm(child)
    return replacement_count

prepare_batchnorms_for_quantization

prepare_batchnorms_for_quantization(model)

Replace frozen BN and set all BN to eval mode with tracking disabled.

Returns the total number of frozen BN modules replaced.

Source code in src/versatil/post_training_compression/preparation/batchnorm.py

def prepare_batchnorms_for_quantization(model: nn.Module) -> int:
    """Replace frozen BN and set all BN to eval mode with tracking disabled.

    Returns the total number of frozen BN modules replaced.
    """
    replacement_count = replace_frozen_batchnorm(model)
    for module in model.modules():
        if isinstance(module, STANDARD_BATCHNORM_TYPES):
            module.eval()
            module.track_running_stats = False
    return replacement_count