# sparselearning package

## Subpackages

- sparselearning.counting
  - sparselearning.counting.helper
  - sparselearning.counting.inference_train_FLOPs
  - sparselearning.counting.micronet_challenge
  - sparselearning.counting.ops
  - sparselearning.counting.print_stats
- sparselearning.funcs
  - sparselearning.funcs.decay
  - sparselearning.funcs.grow
  - sparselearning.funcs.init_scheme
  - sparselearning.funcs.prune
  - sparselearning.funcs.redistribute
- sparselearning.tests
  - sparselearning.tests.test_data
  - sparselearning.tests.test_mask_loading_saving
  - sparselearning.tests.test_struct_sparse
- sparselearning.utils
  - sparselearning.utils.accuracy_helper
  - sparselearning.utils.layer_wise_density
  - sparselearning.utils.model_serialization
  - sparselearning.utils.ops
  - sparselearning.utils.smoothen_value
  - sparselearning.utils.tqdm_logging
  - sparselearning.utils.train_helper
  - sparselearning.utils.warmup_scheduler

## sparselearning.core

Wraps PyTorch model parameters with a boolean mask to simulate unstructured sparsity.

Example usage:

  optimizer = torchoptim.SGD(model.parameters(),lr=args.lr) decay = CosineDecay(args.prune_rate, len(train_loader)*(args.epochs)) mask = Masking(optimizer, prune_rate_decay=decay) model = MyModel() mask.add_module(model)

  # Wrapped optimizer step mask.step

  # Mask update step mask.update_connections()

class `sparselearning.core.``LayerStats`(variance_dict: Dict[str, float] = <factory>, zeros_dict: Dict[str, int] = <factory>, nonzeros_dict: Dict[str, int] = <factory>, removed_dict: Dict[str, int] = <factory>, total_variance: float = 0, total_zero: int = 0, total_nonzero: int = 0, total_removed: int = 0)

  Bases: `object`

  Layer-wise statistics

  `load_state_dict`(*initial_data, **kwargs)
  `nonzeros_dict`: Dict[str, int]
  `removed_dict`: Dict[str, int]
  `state_dict`()
  property `total_density`
  `total_nonzero`: int = 0
  `total_removed`: int = 0
  `total_variance`: float = 0
  `total_zero`: int = 0
  `variance_dict`: Dict[str, float]
  `zeros_dict`: Dict[str, int]

class `sparselearning.core.``Masking`(optimizer: optim, prune_rate_decay: Decay, density: float = 0.1, sparse_init: str = 'random', dense_gradients: bool = False, prune_mode: str = 'magnitude', growth_mode: str = 'momentum', redistribution_mode: str = 'momentum', prune_threshold: float = 0.001, growth_threshold: float = 0.001, growth_increment: float = 0.2, increment: float = 0.2, tolerance: float = 1e-06, mask_dict: Dict[str, Tensor] = <factory>, module: nn.Module = None, mask_step: int = 0)

  Bases: `object`

  Wraps PyTorch model parameters with a sparse mask.

  Creates a mask for each parameter tensor contained in the model. When apply_mask() is called, it applies the sparsity pattern to the parameters.

  Basic usage:

    optimizer = torchoptim.SGD(model.parameters(),lr=args.lr) decay = CosineDecay(args.prune_rate, len(train_loader)*(args.epochs)) mask = Masking(optimizer, prune_rate_decay=decay) model = MyModel() mask.add_module(model)

  Removing layers: Layers can be removed individually, by type, or by partial match of their name.

  > - mask.remove_weight(name) requires an exact name of

  a parameter.

    - mask.remove_weight_partial_name(partial_name=name) removes all parameters that contain the partial name. For example ‘conv’ would remove all layers with ‘conv’ in their name.

    - mask.remove_type(type) removes all layers of a certain type. For example, mask.remove_type(torch.nn.BatchNorm2d) removes all 2D batch norm layers.

  `add_module`(module, lottery_mask_path: Path = None)

    Store dict of parameters to mask

    Parameters

      - module – to mask

      - lottery_mask_path – initialize from an existing model’s mask.

    Returns

  `adjust_prune_rate`()

    Modify prune rate for layers with low sparsity

  `apply_mask`()

    Applies boolean mask to modules

  `apply_mask_gradients`()

    Applies boolean mask to modules’s gradients

  property `avg_inference_FLOPs`
    Returns

      running average of inference FLOPs

    Return type

      float

  `calc_redistributed_densities`()

    Computes layer-wise density given a redistribution scheme.

    Ensures that layer-wise densities are valid (i.e. 0 <= density <= 1).

    Returns

      Layer-wise valid densities.

    Return type

      Dict[str, float]

  property `dense_FLOPs`

    Calculates dense inference FLOPs of the model

    Returns

      dense FLOPs

    Return type

      int

  `dense_gradients` = False
  `density` = 0.1
  `gather_statistics`()

    Gather layer-wise & global stats. Typically performed before each mask update.

  `get_momentum_for_weight`(weight: str) → Tensor

    Return momentum from optimizer (SGD or Adam)

    Parameters

      weight (str) – weight name

    Returns

      Momentum buffer for layer

    Return type

      torch.Tensor

  property `global_prune`
  property `growth_func`
  `growth_increment` = 0.2
  `growth_mode` = 'momentum'
  `growth_threshold` = 0.001
  `increment` = 0.2
  property `inference_FLOPs`

    Calculates dense inference FLOPs of the model

    Returns

      inference FLOPs

    Return type

      float

  `init`(lottery_mask_path: Path)

    Sparsity initialization

    Parameters

      lottery_mask_path (Path) – Mask path, if using Lottery Ticket Hypothesis (Frankle & Carbin 2018).

  `load_state_dict`(*initial_data, **kwargs)
  `mask_step` = 0
  `module` = None
  `print_nonzero_counts`()
  property `prune_func`

    Calls prune func from the registry.

    We use @property, so that it is always synced with prune_mode

  `prune_mode` = 'magnitude'
  property `prune_rate`

    Get prune rate from the decay object

  `prune_threshold` = 0.001
  property `redistribution_func`

    Calls redistribution func from the registry.

    We use @property, so that it is always synced with redistribution_mode

  `redistribution_mode` = 'momentum'
  `remove_type`(nn_type)

    Remove layer by type (eg: nn.Linear, nn.Conv2d, etc.)

    Parameters

      nn_type (nn.Module) – type of layer

  `remove_weight`(name)

    Remove layer by complete name

    Parameters

      name (str) – layer name

  `remove_weight_partial_name`(partial_name: str)

    Remove module by partial name (eg: conv).

    Parameters

      partial_name (str) – partial layer name

  `reset_momentum`()

    Mask momentum buffers

  `sparse_init` = 'random'
  `sparsify`(**kwargs)

    Call sparsity init func (see sparselearning/funcs/init_scheme.py)

  `state_dict`() → Dict
  `step`()

    Performs an optimizer step (i.e, no update to mask topology).

  `to_module_device_`()

    Send to module’s device

  `tolerance` = 1e-06
  `truncate_weights`()

    Perform grow / prune / redistribution step

  `update_connections`()

    Performs a mask update (i.e, update to mask topology).