Source code for neurobench.benchmarks.static_metrics

import torch
import snntorch as snn

# static metrics, only require model




def parameter_count(model):
    """
    Number of parameters in the model.

    Args:
        model: A NeuroBenchModel.
    Returns:
        int: Number of parameters.

    """
    return sum(p.numel() for p in model.__net__().parameters())





def footprint(model):
    """
    Memory footprint of the model.

    Args:
        model: A NeuroBenchModel.
    Returns:
        float: Model size in bytes.

    """
    # Count the number of parameters and multiply by the size of each parameter in bytes
    param_size = 0
    for param in model.__net__().parameters():
        param_size += param.numel() * param.element_size()

    # Count the number of buffers and multiply by the size of each buffer in bytes
    buffer_size = 0
    for buffer in model.__net__().buffers():
        buffer_size += buffer.numel() * buffer.element_size()

    # Return the sum of the parameters and buffers
    return param_size + buffer_size





def connection_sparsity(model):
    """Sparsity of model connections between layers. Based on number of zeros
    in supported layers, other layers are not taken into account in the computation:
    Supported layers:
    Linear
    Conv1d, Conv2d, Conv3d
    RNN, RNNBase, RNNCell
    LSTM, LSTMBase, LSTMCell
    GRU, GRUBase, GRUCell

    Args:
        model: A NeuroBenchModel.
    Returns:
        float: Connection sparsity, rounded to 3 decimals.
    """

    def get_nr_zeros_weights(module):
        """
        Get the number of zeros in a module's weights.

        Args:
            module: A torch.nn.Module.
        Returns:
            int: Number of zeros in the module's weights.

        """
        children = list(module.children())
        regular_layers = (
            torch.nn.Linear,
            torch.nn.Conv1d,
            torch.nn.Conv2d,
            torch.nn.Conv3d,
        )
        recurr_layers = torch.nn.RNNBase
        recurr_cells = torch.nn.RNNCellBase
        if len(children) == 0:  # it is a leaf
            # print(module)
            if isinstance(module, regular_layers):
                count_zeros = torch.sum(module.weight == 0)
                count_weights = module.weight.numel()
                return count_zeros, count_weights

            elif isinstance(module, recurr_layers):
                attribute_names = []
                for i in range(module.num_layers):
                    param_names = ["weight_ih_l{}{}", "weight_hh_l{}{}"]
                    if module.bias:
                        param_names += ["bias_ih_l{}{}", "bias_hh_l{}{}"]
                    if module.proj_size > 0:  # it is lstm
                        param_names += ["weight_hr_l{}{}"]

                    attribute_names += [x.format(i, "") for x in param_names]
                    if module.bidirectional:
                        suffix = "_reverse"
                        attribute_names += [x.format(i, suffix) for x in param_names]

                count_zeros = 0
                count_weights = 0
                for attr in attribute_names:
                    attr_val = getattr(module, attr)
                    count_zeros += torch.sum(attr_val == 0)
                    count_weights += attr_val.numel()

                return count_zeros, count_weights

            elif isinstance(module, recurr_cells):
                attribute_names = ["weight_ih", "weight_hh"]
                if module.bias:
                    attribute_names += ["bias_ih", "bias_hh"]

                count_zeros = 0
                count_weights = 0
                for attr in attribute_names:
                    attr_val = getattr(module, attr)
                    count_zeros += torch.sum(attr_val == 0)
                    count_weights += attr_val.numel()

                return count_zeros, count_weights

            elif isinstance(module, snn.SpikingNeuron):
                return 0, 0  # it is a neuromorphic neuron layer
            else:
                # print('Module type: ', module, 'not found.')
                return 0, 0

        else:
            count_zeros = 0
            count_weights = 0
            for child in children:
                child_zeros, child_weights = get_nr_zeros_weights(child)
                count_zeros += child_zeros
                count_weights += child_weights
            return count_zeros, count_weights

    # Pull the layers from the model's network
    layers = model.__net__().children()
    # For each layer, count where the weights are zero
    count_zeros = 0
    count_weights = 0
    for module in layers:
        zeros, weights = get_nr_zeros_weights(module)
        count_zeros += zeros
        count_weights += weights

    # Return the ratio of zeros to weights, rounded to 4 decimals
    return round((count_zeros / count_weights).item(), 4)