qnn_module.py

"""Sparse Quantized Neural Network torch module."""
from typing import Set, Type

import brevitas.nn as qnn
import numpy
import torch
import torch.nn.utils.prune as pruning
from brevitas.quant.scaled_int import Int8ActPerTensorFloat, Int8WeightPerTensorFloat, IntBias
from torch import nn

from ..common.debugging import assert_true
from ..quantization.qat_quantizers import Int8ActPerTensorPoT, Int8WeightPerTensorPoT


class SparseQuantNeuralNetwork(nn.Module):
    """Sparse Quantized Neural Network.

    This class implements an MLP that is compatible with FHE constraints. The weights and
    activations are quantized to low bit-width and pruning is used to ensure accumulators do not
    surpass an user-provided accumulator bit-width. The number of classes and number of layers
    are specified by the user, as well as the breadth of the network
    """

    # pylint: disable-next=too-many-arguments
    def __init__(
        self,
        input_dim: int,
        n_layers: int,
        n_outputs: int,
        n_hidden_neurons_multiplier: int = 4,
        n_w_bits: int = 4,
        n_a_bits: int = 4,
        # No pruning by default as roundPBS keeps the PBS precision low
        n_accum_bits: int = 32,
        n_prune_neurons_percentage: float = 0.0,
        activation_function: Type = nn.ReLU,
        quant_narrow: bool = False,
        quant_signed: bool = True,
        power_of_two_scaling: bool = True,  # Default to true: use roundPBS to speed up the NNs
    ):
        """Sparse Quantized Neural Network constructor.

        Args:
            input_dim (int): Number of dimensions of the input data.
            n_layers (int): Number of linear layers for this network.
            n_outputs (int): Number of output classes or regression targets.
            n_w_bits (int): Number of weight bits.
            n_a_bits (int): Number of activation and input bits.
            n_accum_bits (int): Maximal allowed bit-width of intermediate accumulators.
            n_hidden_neurons_multiplier (int): The number of neurons on the hidden will be the
                number of dimensions of the input multiplied by `n_hidden_neurons_multiplier`. Note
                that pruning is used to adjust the accumulator size to attempt to keep the maximum
                accumulator bit-width to `n_accum_bits`, meaning that not all hidden layer neurons
                will be active. The default value for `n_hidden_neurons_multiplier` is chosen for
                small dimensions of the input. Reducing this value decreases the FHE inference time
                considerably but also decreases the robustness and accuracy of model training.
            n_prune_neurons_percentage (float): The percentage of neurons to prune in the hidden
                layers. This can be used when setting `n_hidden_neurons_multiplier` with a high
                number (3-4), once good accuracy is obtained, in order to speed up the model in FHE.
            activation_function (Type): The activation function to use in the network
                (e.g., torch.ReLU, torch.SELU, torch.Sigmoid, ...).
            quant_narrow (bool): Whether this network should quantize the values using narrow range
                (e.g a 2-bits signed quantization uses [-1, 0, 1] instead of [-2, -1, 0, 1]).
            quant_signed (bool): Whether this network should quantize the values using signed
                integers.
            power_of_two_scaling (bool): Force quantization scales to be a power of two
                to enable inference speed optimizations. Defaults to True

        Raises:
            ValueError: If the parameters have invalid values or the computed accumulator bit-width
                is zero.
        """

        super().__init__()

        self.features = nn.Sequential()
        in_features = input_dim

        self.n_layers = n_layers

        if n_layers <= 0:
            raise ValueError(
                f"Invalid number of layers: {n_layers}, at least one intermediary layers is needed"
            )

        if n_w_bits <= 0 or n_a_bits <= 0:
            raise ValueError("The weight & activation quantization bit-width cannot be less than 1")

        high_input_bitwidth = False  # power_of_two_scaling and activation_function is nn.ReLU

        for idx in range(n_layers):
            out_features = (
                n_outputs if idx == n_layers - 1 else int(input_dim * n_hidden_neurons_multiplier)
            )

            quant_name = f"quant{idx}"
            quantizer = qnn.QuantIdentity(
                bit_width=8 if high_input_bitwidth else n_a_bits,
                return_quant_tensor=True,
                narrow_range=quant_narrow,
                signed=quant_signed,
                act_quant=Int8ActPerTensorPoT if power_of_two_scaling else Int8ActPerTensorFloat,
            )

            layer_name = f"fc{idx}"
            layer = qnn.QuantLinear(
                in_features,
                out_features,
                True,
                weight_bit_width=n_w_bits,
                bias_quant=IntBias if power_of_two_scaling else None,
                weight_narrow_range=quant_narrow,
                narrow_range=quant_narrow,
                signed=quant_signed,
                weight_quant=Int8WeightPerTensorPoT
                if power_of_two_scaling
                else Int8WeightPerTensorFloat,
            )

            self.features.add_module(quant_name, quantizer)
            self.features.add_module(layer_name, layer)

            if idx < n_layers - 1:
                self.features.add_module(f"act{idx}", activation_function())

            in_features = out_features

        self.n_w_bits = n_w_bits
        self.n_a_bits = n_a_bits
        self.n_accum_bits = n_accum_bits

        # Store input/output dimensions to check they are correct during .fit(X,y).
        # The X passed to .fit must not have different dimensions than the one given in this
        # constructor.
        self.n_outputs = n_outputs
        self.input_dim = input_dim

        self.n_prune_neurons_percentage = n_prune_neurons_percentage

        assert_true(
            self.n_prune_neurons_percentage >= 0 and self.n_prune_neurons_percentage < 1.0,
            "Pruning percentage must be expressed as a fraction between 0 and 1. A value of "
            " zero (0) means pruning is disabled",
        )
        self.pruned_layers: Set[nn.Module] = set()

        self.enable_pruning()

    def max_active_neurons(self) -> int:
        """Compute the maximum number of active (non-zero weight) neurons.

        The computation is done using the quantization parameters passed to the constructor.
        Warning: With the current quantization algorithm (asymmetric) the value returned by this
        function is not guaranteed to ensure FHE compatibility. For some weight distributions,
        weights that are 0 (which are pruned weights) will not be quantized to 0.
        Therefore the total number of active quantized neurons will not be equal to
        max_active_neurons.

        Returns:
            int: The maximum number of active neurons.
        """

        return int(
            numpy.floor(
                (2**self.n_accum_bits - 1) / (2**self.n_w_bits - 1) / (2**self.n_a_bits - 1)
            )
        )

    def make_pruning_permanent(self) -> None:
        """Make the learned pruning permanent in the network."""
        max_neuron_connections = self.max_active_neurons()

        prev_layer_keep_idxs = None
        layer_idx = 0
        # Iterate over all layers that have weights (Linear ones)
        for layer in self.features:
            if not isinstance(layer, (nn.Linear, qnn.QuantLinear)):
                continue

            layer_shape = layer.weight.shape

            # Compute the fan-in, the number of inputs to a neuron, the product of the kernel
            # width x height x in_channels.
            fan_in = numpy.prod(layer_shape[1:])

            # If this is a layer that should be pruned and is currently being pruned, make the
            # pruning permanent. This is done by multiplying the pruning mask tensor with the
            # weights and storing the result in the weight member.
            if layer in self.pruned_layers and fan_in > max_neuron_connections:
                pruning.remove(layer, "weight")
                self.pruned_layers.remove(layer)

            if self.n_prune_neurons_percentage > 0.0:
                weights = layer.weight.detach().numpy()

                # Pruning all layers except the last one, but for the last one
                # we still need to remove synapses of the previous layer's pruned neurons
                if layer_idx < self.n_layers - 1:
                    # Once pruning is disabled, the weights of some neurons become 0
                    # We need to find those neurons (columns in the weight matrix).
                    # Testing for floats equal to 0 is done using an epsilon
                    neurons_removed_idx = numpy.where(numpy.sum(numpy.abs(weights), axis=1) < 0.001)
                    idx = numpy.arange(weights.shape[0])
                    keep_idxs = numpy.setdiff1d(idx, neurons_removed_idx)

                    # Remove the pruning hooks on this layer
                    if layer in self.pruned_layers:
                        pruning.remove(layer, "weight")
                        self.pruned_layers.remove(layer)
                else:
                    keep_idxs = numpy.arange(weights.shape[0])

                # Now we take the indices of the neurons kept for the previous layer
                # If this is the first layer all neurons are kept
                if prev_layer_keep_idxs is None:
                    prev_layer_keep_idxs = numpy.arange(weights.shape[1])

                # Remove the pruned neurons and the weights/synapses
                # that apply to neurons removed in the previous layer
                orig_weight = layer.weight.data.clone()
                transform_weight = orig_weight[keep_idxs]
                transform_weight = transform_weight[:, prev_layer_keep_idxs]

                # Replace the weight matrix of the current layer
                layer.weight = torch.nn.Parameter(transform_weight)

                # Eliminate the biases of the neurons that were removed in this layer
                if layer.bias is not None:
                    orig_bias = layer.bias.data.clone()
                    transform_bias = orig_bias[keep_idxs]
                    layer.bias = torch.nn.Parameter(transform_bias)

                # Save the indices of the neurons removed in this layer to
                # remove synapses in the next layer
                prev_layer_keep_idxs = keep_idxs

                # Check that pruning mask introduced by torch pruning was removed
                assert not hasattr(layer, "weight_mask")

            layer_idx += 1

        assert_true(
            len(self.pruned_layers) == 0,
            "Making pruning permanent did not remove pruning on all layers",
        )

        assert_true(
            layer_idx == self.n_layers,
            "Not all layers in the network were examined as candidates for pruning",
        )

    def enable_pruning(self) -> None:
        """Enable pruning in the network. Pruning must be made permanent to recover pruned weights.

        Raises:
            ValueError: If the quantization parameters are invalid.
        """
        max_neuron_connections = self.max_active_neurons()

        if max_neuron_connections == 0:
            raise ValueError(
                "The maximum accumulator bit-width is too low "
                "for the quantization parameters requested. No neurons would be created in the "
                "requested configuration"
            )

        # Iterate over all layers that have weights (Linear ones)
        layer_idx = 0
        for layer in self.features:
            if not isinstance(layer, (nn.Linear, qnn.QuantLinear)):
                continue

            layer_shape = layer.weight.shape

            # Compute the fan-in, the number of inputs to a neuron, and the fan-out, the number of
            # neurons in the current layer.
            # The fan-in is the product of the kernel width x height x in_channels while the fan-out
            # is out_channels
            fan_in = numpy.prod(layer_shape[1:])
            fan_out = layer_shape[0]

            # To satisfy accumulator bit-width constraints each dot-product between an input line
            # and weight column must not exceed n_accum_bits bits. We thus prune the layer to have
            # at most max_neuron_connections non-zero weights
            if fan_in > max_neuron_connections and layer not in self.pruned_layers:
                pruning.l1_unstructured(
                    layer, "weight", (fan_in - max_neuron_connections) * fan_out
                )
                self.pruned_layers.add(layer)

            # If pruning is enabled, which is generally the case during training, a
            # forward hook is added in order to create a mask tensor (made of 0 or 1) that will
            # be multiplied with the weights during the forward pass.
            # This is done for all layers except the last one, which outputs the prediction
            if layer_idx < self.n_layers - 1 and self.n_prune_neurons_percentage > 0.0:
                # Use L2-norm structured pruning, using the torch ln_structured
                # function, with norm=2 and axis=0 (output/neuron axis)
                pruning.ln_structured(layer, "weight", self.n_prune_neurons_percentage, 2, 0)
                self.pruned_layers.add(layer)

            # Note this is counting only Linear layers
            layer_idx += 1

        assert_true(
            layer_idx == self.n_layers,
            "Not all layers in the network were examined as candidates for pruning",
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass.

        Args:
            x (torch.Tensor): network input

        Returns:
            x (torch.Tensor): network prediction
        """
        return self.features(x)