sparselearning.counting

sparselearning.counting.helper

sparselearning.counting.helper.get_pre_activations_dict(net: nn.Module, input_tensor: Tensor) → Dict[str, Tensor]

Find pre-activation dict for every possible module in network

Parameters

  • net (nn.Module) – Pytorch model

  • input_tensor (Tensor) – input tensor, supports only single input

Returns

dictionary mapping layers to pre-activations

Return type

Dict[str, Tensor]

sparselearning.counting.inference_train_FLOPs

sparselearning.counting.inference_train_FLOPs.Pruning_inference_FLOPs(dense_FLOPs: int, decay: sparselearning.funcs.decay.MagnitudePruneDecay, total_steps: int = 87891) → float

Inference FLOPs for Iterative Pruning, Zhu and Gupta 2018. Note, assumes FLOPs propto average sparsity, which is approximately true in practice.

For our report, we accurately calculate train FLOPs by evaluating FLOPs during each pruning iteration.

Parameters
Returns

Pruning inference FLOPs

Return type

float

sparselearning.counting.inference_train_FLOPs.Pruning_train_FLOPs(dense_FLOPs: int, decay: sparselearning.funcs.decay.MagnitudePruneDecay, total_steps: int = 87891) → float

Train FLOPs for Iterative Pruning, Zhu and Gupta 2018. Note, assumes FLOPs propto average sparsity, which is approximately true in practice.

For our report, we accurately calculate train FLOPs by evaluating FLOPs during each pruning iteration.

Parameters
Returns

Pruning train FLOPs

Return type

float

sparselearning.counting.inference_train_FLOPs.RigL_train_FLOPs(sparse_FLOPs: int, dense_FLOPs: int, mask_interval: int = 100) → float

Train FLOPs for Rigging the Lottery (RigL), Evci et al. 2020.

Parameters

  • sparse_FLOPs (int) – FLOPs consumed for sparse model’s forward pass

  • dense_FLOPs (int) – FLOPs consumed for dense model’s forward pass

  • mask_interval (int) – Mask update interval

Returns

RigL train FLOPs

Return type

float

sparselearning.counting.inference_train_FLOPs.SET_train_FLOPs(sparse_FLOPs: int, dense_FLOPs: int, mask_interval: int = 100) → int

Train FLOPs for Sparse Evolutionary Training (SET), Mocanu et al. 2018.

Parameters

  • sparse_FLOPs (int) – FLOPs consumed for sparse model’s forward pass

  • dense_FLOPs (int) – FLOPs consumed for dense model’s forward pass

  • mask_interval (int) – Mask update interval

Returns

SET train FLOPs

Return type

int

sparselearning.counting.inference_train_FLOPs.SNFS_train_FLOPs(sparse_FLOPs: int, dense_FLOPs: int, mask_interval: int = 100) → int

Train FLOPs for Sparse Networks from Scratch (SNFS), Dettmers et al. 2020.

Parameters

  • sparse_FLOPs (int) – FLOPs consumed for sparse model’s forward pass

  • dense_FLOPs (int) – FLOPs consumed for dense model’s forward pass

  • mask_interval (int) – Mask update interval

Returns

SNFS train FLOPs

Return type

int

sparselearning.counting.inference_train_FLOPs.model_inference_FLOPs(sparse_init: str = 'random', density: float = 0.2, model_name: str = 'wrn-22-2', input_size: Tuple = (1, 3, 32, 32)) → int

Obtain inference FLOPs for a model.

Only for models trained with a constant FLOP sparsifying technique. eg: SNFS, Pruning are not supported here.

Parameters

  • sparse_init (str) – Initialization scheme used (Random / ER / ERK)

  • density (float) – Overall parameter density (non-zero / capacity)

  • model_name (str) – model to use (WideResNet-22-2 or ResNet-50)

  • input_size (Tuple) – shape of input tensor

Returns

Return type

sparselearning.counting.inference_train_FLOPs.resnet50_FLOPs(sparse_init: str = 'random', density: float = 0.2, *, model_name: str = 'resnet50', input_size: Tuple = (1, 3, 32, 32)) → int

Obtain inference FLOPs for a model.

Only for models trained with a constant FLOP sparsifying technique. eg: SNFS, Pruning are not supported here.

Parameters

  • sparse_init (str) – Initialization scheme used (Random / ER / ERK)

  • density (float) – Overall parameter density (non-zero / capacity)

  • model_name (str) – model to use (WideResNet-22-2 or ResNet-50)

  • input_size (Tuple) – shape of input tensor

Returns

Return type

sparselearning.counting.inference_train_FLOPs.wrn_22_2_FLOPs(sparse_init: str = 'random', density: float = 0.2, *, model_name: str = 'wrn-22-2', input_size: Tuple = (1, 3, 32, 32)) → int

Obtain inference FLOPs for a model.

Only for models trained with a constant FLOP sparsifying technique. eg: SNFS, Pruning are not supported here.

Parameters

  • sparse_init (str) – Initialization scheme used (Random / ER / ERK)

  • density (float) – Overall parameter density (non-zero / capacity)

  • model_name (str) – model to use (WideResNet-22-2 or ResNet-50)

  • input_size (Tuple) – shape of input tensor

Returns

Return type

sparselearning.counting.micronet_challenge

Code taken from: https://github.com/google-research/google-research/blob/master/micronet_challenge/counting.py

This module defines an API for counting parameters and operations.

## Defining the Operation Count API - input_size is an int, since square image assumed. - strides is a tuple, but assumed to have same stride in both dimensions. - Supported paddings are same’ and `valid. - use_bias is boolean. - activation is one of the following relu, swish, sigmoid, None - kernel_shapes for Conv2D dimensions must be in the following order:

k_size, k_size, c_in, c_out

  • kernel_shapes for FullyConnected dimensions must be in the following order: c_in, c_out

  • kernel_shapes for DepthWiseConv2D dimensions must be like the following: k_size, k_size, c_in==c_out, 1

class sparselearning.counting.micronet_challenge.Add(input_size, n_channels)

Bases: tuple

Operation definitions for elementwise multiplication and addition.

Attributes:

kernel_shape: list, of length 2. Shape of the weight matrix. use_bias: bool, if true a bias term is added to the output. activation: str, type of activation applied to the output.

input_size

Alias for field number 0

n_channels

Alias for field number 1

class sparselearning.counting.micronet_challenge.Conv2D(input_size, kernel_shape, strides, padding, use_bias, activation)

Bases: tuple

Operation definition for 2D depthwise convolution.

Only difference compared to Conv2D is the kernel_shape[3] = 1.

activation

Alias for field number 5

input_size

Alias for field number 0

kernel_shape

Alias for field number 1

padding

Alias for field number 3

strides

Alias for field number 2

use_bias

Alias for field number 4

class sparselearning.counting.micronet_challenge.DepthWiseConv2D(input_size, kernel_shape, strides, padding, use_bias, activation)

Bases: tuple

Operation definition for Global Average Pooling.

Attributes:

input_size: int, Dimensions of the input image (square assumed). n_channels: int, Number of output dimensions.

activation

Alias for field number 5

input_size

Alias for field number 0

kernel_shape

Alias for field number 1

padding

Alias for field number 3

strides

Alias for field number 2

use_bias

Alias for field number 4

class sparselearning.counting.micronet_challenge.FullyConnected(kernel_shape, use_bias, activation)

Bases: tuple

activation

Alias for field number 2

kernel_shape

Alias for field number 0

use_bias

Alias for field number 1

class sparselearning.counting.micronet_challenge.GlobalAvg(input_size, n_channels)

Bases: tuple

Operation definitions for elementwise multiplication and addition.

Attributes:

input_size: int, Dimensions of the input image (square assumed). n_channels: int, Number of output dimensions.

input_size

Alias for field number 0

n_channels

Alias for field number 1

class sparselearning.counting.micronet_challenge.MicroNetCounter(all_ops, add_bits_base=32, mul_bits_base=32)

Bases: object

Counts operations using given information.

print_summary(sparsity, param_bits, add_bits, mul_bits, summarize_blocks=True)

Prints all operations with given options.

Args:

sparsity: float, between 0,1 defines how sparse each parametric layer is. param_bits: int, bits in which parameters are stored. add_bits: float, number of bits used for accumulator. mul_bits: float, number of bits inputs represented for multiplication. summarize_blocks: bool, if True counts within a block are aggregated and

reported in a single line.

process_counts(total_params, total_mults, total_adds, mul_bits, add_bits)
class sparselearning.counting.micronet_challenge.Scale(input_size, n_channels)

Bases: tuple

input_size

Alias for field number 0

n_channels

Alias for field number 1

sparselearning.counting.micronet_challenge.count_ops(op, sparsity, param_bits)

Given a operation class returns the flop and parameter statistics.

Args:

op: namedtuple, operation definition. sparsity: float, sparsity of parameterized operations. Sparsity only effects

Conv and FC layers; since activations are dense.

param_bits: int, number of bits required to represent a parameter.

Returns:

param_count: number of bits required to store parameters n_mults: number of multiplications made per input sample. n_adds: number of multiplications made per input sample.

sparselearning.counting.micronet_challenge.get_conv_output_size(image_size, filter_size, padding, stride)

Calculates the output size of convolution.

The input, filter and the strides are assumed to be square. Arguments:

image_size: int, Dimensions of the input image (square assumed). filter_size: int, Dimensions of the kernel (square assumed). padding: str, padding added to the input image. ‘same’ or ‘valid’ stride: int, stride with which the kernel is applied (square assumed).

Returns:

int, output size.

sparselearning.counting.micronet_challenge.get_flops_per_activation(activation)

Returns the number of multiplication ands additions of an activation.

Args:

activation: str, type of activation applied to the output.

Returns:

n_muls, n_adds

sparselearning.counting.micronet_challenge.get_info(op)

Given an op extracts some common information.

sparselearning.counting.micronet_challenge.get_sparse_size(tensor_shape, param_bits, sparsity)

Given a tensor shape returns #bits required to store the tensor sparse.

If sparsity is greater than 0, we do have to store a bit mask to represent sparsity. Args:

tensor_shape: list<int>, shape of the tensor param_bits: int, number of bits the elements of the tensor represented in. sparsity: float, sparsity level. 0 means dense.

Returns:

int, number of bits required to represented the tensor in sparse format.

sparselearning.counting.ops

sparselearning.counting.ops.get_inference_FLOPs(masking: Masking, input_tensor: Tensor, param_size: int = 32) → int

Returns total FLOPs consumed for a forward pass (inference FLOPs). Assumes support for sparse convolutions, sparse dense layers.

Parameters

  • masking (Masking) – Masking instance, a wrapper on boolean masks

  • input_tensor (torch.Tensor) – Input to model, single input supported

  • param_size (int) – bits used for floating point operations (default 32)

Returns

total FLOPs consumed for a forward pass

Return type

int

sparselearning.counting.print_stats

sparselearning.counting.print_stats.print_stats(model_name: str, input_size: Tuple = (1, 3, 32, 32))

Print FLOP statistics for (Dense, Pruning, RigL, SET, SNFS) models

Parameters

  • model_name (str) – Model to use (wrn-22-2 or resnet-50)

  • input_size (Tuple[int]) – Shape of input tensor, single input only