base

base

Abstract base class for weight pruning strategies.

BasePruner

Bases: ABC

Base interface for all pruning strategies.

prune abstractmethod

prune(module)

Apply pruning to module.

Parameters:

Name	Type	Description	Default
module	Module	Neural network module to prune.	required

Returns:

Type	Description
tuple[int, int]	Tuple of (total_parameters, zero_parameters).

Source code in src/versatil/post_training_compression/pruning/base.py

@abc.abstractmethod
def prune(self, module: nn.Module) -> tuple[int, int]:
    """Apply pruning to module.

    Args:
        module: Neural network module to prune.

    Returns:
        Tuple of (total_parameters, zero_parameters).
    """

compute_sparsity staticmethod

compute_sparsity(module)

Count total and zero parameters in module.

Parameters:

Name	Type	Description	Default
module	Module	Neural network module to inspect.	required

Returns:

Type	Description
tuple[int, int]	Tuple of (total_parameters, zero_parameters).

Source code in src/versatil/post_training_compression/pruning/base.py

@staticmethod
def compute_sparsity(module: nn.Module) -> tuple[int, int]:
    """Count total and zero parameters in module.

    Args:
        module: Neural network module to inspect.

    Returns:
        Tuple of (total_parameters, zero_parameters).
    """
    total_parameters = 0
    zero_parameters = 0
    for parameter in module.parameters():
        total_parameters += parameter.numel()
        zero_parameters += int(torch.sum(parameter == 0).item())
    return total_parameters, zero_parameters