ANN-to-SNN Conversion

Convert trained PyTorch ANNs to rate-coded spiking neural networks.

Contract

The conversion package is an optional PyTorch surface. The base package can be imported without PyTorch; resolving convert, ConvertedSNN, or QCFSActivation requires a PyTorch-capable environment.

• convert(model, calibration_data=None, T=None, percentile=99.9) extracts Linear and Conv2d weights and selects a conversion route from the model's activations, returning a deterministic ConvertedSNN:
• Threshold-balancing route (ReLU models, Diehl et al. 2015) — calibrates activation thresholds from ReLU layers when calibration data is supplied, starts each IF neuron from rest, and defaults T to 16.
• QCFS route (QCFS-trained models, Bu et al. 2022) — takes each QCFSActivation's learned threshold directly, pre-loads each IF neuron to theta / 2 (the optimal shift that cancels the quantisation bias), ignores calibration data, and adopts the layers' trained step budget when T is left unset.
• ConvertedSNN.run(x) rate-codes NumPy input with a fixed RNG seed and returns output spike counts for one vector or a batch. initial_membrane_fraction controls the per-layer membrane pre-load (0.0 rest, 0.5 QCFS shift).
• ConvertedSNN.classify(x) returns the argmax class index from output spike counts.
• QCFSActivation replaces ReLU during conversion-aware training by clipping activations to [0, theta] and quantising them to T + 1 spike-rate levels with a straight-through gradient.
• replace_relu_with_qcfs(model, T=8, theta=1.0, learn_theta=True) substitutes every ReLU/ReLU6 in a model (recursing through submodules) with a QCFSActivation, preparing the network for QCFS conversion-aware fine-tuning.

Verification

The public conversion files are covered by the scoped NumPy-docstring policy:

• src/sc_neurocore/conversion/__init__.py
• src/sc_neurocore/conversion/ann_to_snn.py
• src/sc_neurocore/conversion/qcfs.py

Focused production tests live in tests/test_conversion.py and tests/test_conversion_ann_snn.py. They exercise real PyTorch modules, the threshold-balancing and QCFS conversion routes, ConvertedSNN.run, ConvertedSNN.classify, the membrane shift, the ReLU→QCFS substitution helper, QCFS range and gradient behaviour, and the layer-extraction contract.

Converter

sc_neurocore.conversion.ann_to_snn

Convert trained PyTorch ANNs to rate-coded spiking neural networks.

Two conversion routes are supported, selected automatically from the activations present in the source model:

Threshold-balancing route (ReLU models, Diehl et al. 2015) Replaces ReLU activations with integrate-and-fire (IF) neurons and rescales weights so that the calibrated maximum activation maps onto the firing threshold. A spike train of rate a / theta over T timesteps approximates the ANN activation a.

QCFS route (QCFS-trained models, Bu et al. 2022) When the model carries :class:~sc_neurocore.conversion.qcfs.QCFSActivation layers, their learned per-layer thresholds become the IF thresholds directly (no calibration pass) and each IF neuron is initialised to a membrane potential of theta / 2 — the optimal shift that cancels the quantisation flooring bias. A QCFS-trained ANN then converts to an SNN with near-zero accuracy loss at the matching timestep budget.

Pipeline

1. Extract weights and biases from the PyTorch model.
2. Derive per-layer thresholds — from QCFS layers when present, otherwise from calibration activation statistics.
3. Normalise weights so each threshold maps onto unity.
4. Build an IF-neuron SNN that reproduces the ANN output as spike counts over T timesteps, with the QCFS membrane shift applied when converting a QCFS-trained model.

References

Diehl et al. 2015 — "Fast-classifying, high-accuracy spiking deep networks through weight and threshold balancing". Bu et al. 2022 — "Optimal ANN-SNN Conversion for High-accuracy and Ultra-low-latency Spiking Neural Networks" (ICLR).

ConvertedSNN dataclass

Rate-coded SNN converted from an ANN.

Attributes

weights : list of ndarray Per-layer weight matrices. biases : list of ndarray or None Per-layer biases (None if absent). thresholds : list of float Per-layer firing thresholds after normalization. T : int Number of simulation timesteps. initial_membrane_fraction : float Fraction of each layer's threshold pre-loaded into the IF membrane potential before the first timestep. 0.0 reproduces the threshold-balancing route; 0.5 applies the QCFS optimal shift (Bu et al. 2022) that cancels the quantisation flooring bias. n_layers : int Number of layers.

Source code in src/sc_neurocore/conversion/ann_to_snn.py

@dataclass
class ConvertedSNN:
    """Rate-coded SNN converted from an ANN.

    Attributes
    ----------
    weights : list of ndarray
        Per-layer weight matrices.
    biases : list of ndarray or None
        Per-layer biases (None if absent).
    thresholds : list of float
        Per-layer firing thresholds after normalization.
    T : int
        Number of simulation timesteps.
    initial_membrane_fraction : float
        Fraction of each layer's threshold pre-loaded into the IF membrane
        potential before the first timestep. ``0.0`` reproduces the
        threshold-balancing route; ``0.5`` applies the QCFS optimal shift
        (Bu et al. 2022) that cancels the quantisation flooring bias.
    n_layers : int
        Number of layers.
    """

    weights: list[np.ndarray[Any, Any]]
    biases: list[np.ndarray[Any, Any] | None]
    thresholds: list[float]
    T: int
    initial_membrane_fraction: float = 0.0
    n_layers: int = field(init=False)

    def __post_init__(self) -> None:
        """Derive the layer count from the converted weight stack."""
        self.n_layers = len(self.weights)

    def run(self, x: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
        """Run the converted SNN for T timesteps on input x.

        Parameters
        ----------
        x : ndarray of shape (n_input,) or (batch, n_input)
            Input values in [0, 1]. Converted to Poisson spike trains.

        Returns
        -------
        ndarray of shape (n_output,) or (batch, n_output)
            Output spike counts over T timesteps (unnormalized).
        """
        squeeze = x.ndim == 1
        if squeeze:
            x = x[np.newaxis]

        batch = x.shape[0]
        rng = np.random.RandomState(42)

        # Initialize membrane voltages. The QCFS route pre-loads theta/2 per
        # layer (initial_membrane_fraction == 0.5); the threshold-balancing
        # route starts from rest (0.0).
        voltages = [
            np.full((batch, w.shape[0]), self.initial_membrane_fraction * theta)
            for w, theta in zip(self.weights, self.thresholds)
        ]
        spike_counts = np.zeros((batch, self.weights[-1].shape[0]))

        for t in range(self.T):
            # Rate-code input: spike with probability proportional to x
            input_spikes = (rng.random(x.shape) < x).astype(np.float64)

            layer_input = input_spikes
            for i, (w, b, theta) in enumerate(zip(self.weights, self.biases, self.thresholds)):
                current = layer_input @ w.T
                if b is not None:
                    current += b / self.T
                voltages[i] += current
                spikes = (voltages[i] >= theta).astype(np.float64)
                voltages[i] -= spikes * theta
                layer_input = spikes

                if i == self.n_layers - 1:
                    spike_counts += spikes

        if squeeze:
            spike_counts = spike_counts[0]
        return spike_counts

    def classify(self, x: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
        """Run SNN and return predicted class indices."""
        counts = self.run(x)
        predictions: np.ndarray[Any, Any] = np.argmax(counts, axis=-1)
        return predictions

__post_init__()

Derive the layer count from the converted weight stack.

Source code in src/sc_neurocore/conversion/ann_to_snn.py

def __post_init__(self) -> None:
    """Derive the layer count from the converted weight stack."""
    self.n_layers = len(self.weights)

run(x)

Run the converted SNN for T timesteps on input x.

Parameters

x : ndarray of shape (n_input,) or (batch, n_input) Input values in [0, 1]. Converted to Poisson spike trains.

Returns

ndarray of shape (n_output,) or (batch, n_output) Output spike counts over T timesteps (unnormalized).

Source code in src/sc_neurocore/conversion/ann_to_snn.py

def run(self, x: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
    """Run the converted SNN for T timesteps on input x.

    Parameters
    ----------
    x : ndarray of shape (n_input,) or (batch, n_input)
        Input values in [0, 1]. Converted to Poisson spike trains.

    Returns
    -------
    ndarray of shape (n_output,) or (batch, n_output)
        Output spike counts over T timesteps (unnormalized).
    """
    squeeze = x.ndim == 1
    if squeeze:
        x = x[np.newaxis]

    batch = x.shape[0]
    rng = np.random.RandomState(42)

    # Initialize membrane voltages. The QCFS route pre-loads theta/2 per
    # layer (initial_membrane_fraction == 0.5); the threshold-balancing
    # route starts from rest (0.0).
    voltages = [
        np.full((batch, w.shape[0]), self.initial_membrane_fraction * theta)
        for w, theta in zip(self.weights, self.thresholds)
    ]
    spike_counts = np.zeros((batch, self.weights[-1].shape[0]))

    for t in range(self.T):
        # Rate-code input: spike with probability proportional to x
        input_spikes = (rng.random(x.shape) < x).astype(np.float64)

        layer_input = input_spikes
        for i, (w, b, theta) in enumerate(zip(self.weights, self.biases, self.thresholds)):
            current = layer_input @ w.T
            if b is not None:
                current += b / self.T
            voltages[i] += current
            spikes = (voltages[i] >= theta).astype(np.float64)
            voltages[i] -= spikes * theta
            layer_input = spikes

            if i == self.n_layers - 1:
                spike_counts += spikes

    if squeeze:
        spike_counts = spike_counts[0]
    return spike_counts

classify(x)

Run SNN and return predicted class indices.

Source code in src/sc_neurocore/conversion/ann_to_snn.py

def classify(self, x: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
    """Run SNN and return predicted class indices."""
    counts = self.run(x)
    predictions: np.ndarray[Any, Any] = np.argmax(counts, axis=-1)
    return predictions

convert(model, calibration_data=None, T=None, percentile=99.9)

Convert a trained PyTorch ANN to a rate-coded SNN.

The conversion route is selected from the model's activations: a model carrying :class:QCFSActivation layers takes the QCFS route (learned thresholds, theta / 2 membrane shift, no calibration); any other model takes the threshold-balancing route (calibrated or unit thresholds, rest-state membrane).

Parameters

model : nn.Module Trained PyTorch model with Linear/Conv2d layers and either ReLU or QCFS activations. calibration_data : Tensor, optional Sample input batch for threshold calibration on the ReLU route. If None, the ReLU route uses a default threshold of 1.0 per layer. Ignored on the QCFS route, whose thresholds are already learned. T : int, optional Number of simulation timesteps (higher = more accurate, slower). If None, the QCFS route adopts the layers' trained step budget and the ReLU route defaults to 16. percentile : float Activation percentile for threshold normalization on the ReLU route.

Returns

ConvertedSNN Converted spiking network ready to run.

Source code in src/sc_neurocore/conversion/ann_to_snn.py

def convert(
    model: object,
    calibration_data: object = None,
    T: int | None = None,
    percentile: float = 99.9,
) -> ConvertedSNN:
    """Convert a trained PyTorch ANN to a rate-coded SNN.

    The conversion route is selected from the model's activations: a model
    carrying :class:`QCFSActivation` layers takes the QCFS route (learned
    thresholds, ``theta / 2`` membrane shift, no calibration); any other
    model takes the threshold-balancing route (calibrated or unit
    thresholds, rest-state membrane).

    Parameters
    ----------
    model : nn.Module
        Trained PyTorch model with Linear/Conv2d layers and either ReLU or
        QCFS activations.
    calibration_data : Tensor, optional
        Sample input batch for threshold calibration on the ReLU route. If
        None, the ReLU route uses a default threshold of 1.0 per layer.
        Ignored on the QCFS route, whose thresholds are already learned.
    T : int, optional
        Number of simulation timesteps (higher = more accurate, slower). If
        None, the QCFS route adopts the layers' trained step budget and the
        ReLU route defaults to 16.
    percentile : float
        Activation percentile for threshold normalization on the ReLU route.

    Returns
    -------
    ConvertedSNN
        Converted spiking network ready to run.
    """
    if not HAS_TORCH:
        raise ImportError("PyTorch required for ANN-to-SNN conversion")

    layers = _extract_layers(model)
    if not layers:
        raise ValueError("No Linear/Conv2d layers found in model")

    weights = [w for w, _ in layers]
    biases = [b for _, b in layers]

    qcfs_layers = _extract_qcfs_layers(model)
    if qcfs_layers:
        # QCFS route (Bu et al. 2022): the learned per-layer theta is the
        # threshold and theta/2 is the optimal initial membrane potential.
        thresholds = [theta for theta, _ in qcfs_layers]
        # Pad if fewer QCFS layers than weight layers (e.g. no output QCFS).
        while len(thresholds) < len(weights):
            thresholds.append(1.0)
        initial_membrane_fraction = 0.5
        if T is None:
            T = qcfs_layers[0][1]
    else:
        # Threshold-balancing route (Diehl et al. 2015).
        initial_membrane_fraction = 0.0
        if calibration_data is not None:
            max_acts = _compute_max_activations(
                model, cast(torch.Tensor, calibration_data), percentile
            )
            # Pad if fewer ReLUs than Linear layers
            while len(max_acts) < len(weights):
                max_acts.append(1.0)
            thresholds = max_acts
        else:
            thresholds = [1.0] * len(weights)
        if T is None:
            T = 16

    # Normalize weights: scale so that each threshold maps to unity.
    normalized_weights = []
    prev_scale = 1.0
    for i, (w, theta) in enumerate(zip(weights, thresholds)):
        scale = theta / prev_scale if i > 0 else theta
        normalized_weights.append(w / scale)
        prev_scale = theta

    return ConvertedSNN(
        weights=normalized_weights,
        biases=biases,
        thresholds=[1.0] * len(weights),
        T=T,
        initial_membrane_fraction=initial_membrane_fraction,
    )

replace_relu_with_qcfs(model, T=8, theta=1.0, learn_theta=True)

Swap every ReLU/ReLU6 in a model for a QCFS activation, in place.

This prepares a trained or fresh ANN for conversion-aware fine-tuning: after substitution the network is retrained for a few epochs so the QCFS thresholds settle, after which :func:convert produces a near-lossless SNN (Bu et al. 2022).

Parameters

model : nn.Module Model whose ReLU/ReLU6 activations are replaced. Mutated in place, recursing through every submodule. T : int Quantisation step budget for each inserted QCFS layer. theta : float Initial firing threshold for each inserted QCFS layer. learn_theta : bool Whether each inserted threshold is a trainable parameter (the QCFS fine-tuning default).

Returns

nn.Module The same model instance, returned for chaining.

Source code in src/sc_neurocore/conversion/ann_to_snn.py

def replace_relu_with_qcfs(
    model: Any,
    T: int = 8,
    theta: float = 1.0,
    learn_theta: bool = True,
) -> Any:
    """Swap every ReLU/ReLU6 in a model for a QCFS activation, in place.

    This prepares a trained or fresh ANN for conversion-aware fine-tuning:
    after substitution the network is retrained for a few epochs so the QCFS
    thresholds settle, after which :func:`convert` produces a near-lossless
    SNN (Bu et al. 2022).

    Parameters
    ----------
    model : nn.Module
        Model whose ReLU/ReLU6 activations are replaced. Mutated in place,
        recursing through every submodule.
    T : int
        Quantisation step budget for each inserted QCFS layer.
    theta : float
        Initial firing threshold for each inserted QCFS layer.
    learn_theta : bool
        Whether each inserted threshold is a trainable parameter (the QCFS
        fine-tuning default).

    Returns
    -------
    nn.Module
        The same ``model`` instance, returned for chaining.
    """
    if not HAS_TORCH:
        raise ImportError("PyTorch required for ANN-to-SNN conversion")

    for name, child in model.named_children():
        if isinstance(child, (nn.ReLU, nn.ReLU6)):
            setattr(model, name, QCFSActivation(T=T, theta=theta, learn_theta=learn_theta))
        else:
            replace_relu_with_qcfs(child, T=T, theta=theta, learn_theta=learn_theta)
    return model

QCFS Activation

sc_neurocore.conversion.qcfs

QCFS (Quantization-Clip-Floor-Shift) activation function.

Replaces ReLU in the ANN during conversion-aware training or post-hoc conversion. QCFS approximates the rate-coded SNN firing rate as a quantized step function, minimizing conversion error.

Reference: Bu et al. 2022 — "Optimal ANN-SNN Conversion for High-accuracy and Ultra-low-latency Spiking Neural Networks"

QCFSActivation

Bases: Module

QCFS activation: quantized clip-floor-shift ReLU replacement.

For T timesteps and threshold theta

QCFS(x) = clip(floor(x * T / theta + 0.5), 0, T) * theta / T

This quantizes activations to T+1 levels in [0, theta], matching the achievable spike rates of an IF neuron over T timesteps.

Parameters

T : int Number of simulation timesteps. theta : float Firing threshold (default 1.0). learn_theta : bool Make threshold trainable (default False).

Source code in src/sc_neurocore/conversion/qcfs.py

class QCFSActivation(nn.Module):
    """QCFS activation: quantized clip-floor-shift ReLU replacement.

    For T timesteps and threshold theta:
        QCFS(x) = clip(floor(x * T / theta + 0.5), 0, T) * theta / T

    This quantizes activations to T+1 levels in [0, theta], matching
    the achievable spike rates of an IF neuron over T timesteps.

    Parameters
    ----------
    T : int
        Number of simulation timesteps.
    theta : float
        Firing threshold (default 1.0).
    learn_theta : bool
        Make threshold trainable (default False).
    """

    def __init__(self, T: int = 8, theta: float = 1.0, learn_theta: bool = False) -> None:
        super().__init__()
        self.T = T
        if learn_theta:
            self.theta = nn.Parameter(torch.tensor(theta))
        else:
            self.register_buffer("theta", torch.tensor(theta))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Quantise activations to the spike-rate grid with a straight-through gradient.

        Parameters
        ----------
        x : torch.Tensor
            ANN activation tensor to clip and quantise into ``T + 1`` rate levels.

        Returns
        -------
        torch.Tensor
            Tensor with values clipped to ``[0, theta]`` and quantised to the
            finite-timestep spike-rate lattice.
        """
        scaled = x * self.T / self.theta + 0.5
        # STE: floor in forward, pass gradient straight through
        quantized = scaled.floor() - (scaled.floor() - scaled).detach()
        clipped = quantized.clamp(0, self.T)
        out: torch.Tensor = clipped * self.theta / self.T
        return out

    def extra_repr(self) -> str:
        """Return the compact PyTorch module representation."""
        return f"T={self.T}, theta={self.theta.item():.2f}"

forward(x)

Quantise activations to the spike-rate grid with a straight-through gradient.

Parameters

x : torch.Tensor ANN activation tensor to clip and quantise into T + 1 rate levels.

Returns

torch.Tensor Tensor with values clipped to [0, theta] and quantised to the finite-timestep spike-rate lattice.

Source code in src/sc_neurocore/conversion/qcfs.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Quantise activations to the spike-rate grid with a straight-through gradient.

    Parameters
    ----------
    x : torch.Tensor
        ANN activation tensor to clip and quantise into ``T + 1`` rate levels.

    Returns
    -------
    torch.Tensor
        Tensor with values clipped to ``[0, theta]`` and quantised to the
        finite-timestep spike-rate lattice.
    """
    scaled = x * self.T / self.theta + 0.5
    # STE: floor in forward, pass gradient straight through
    quantized = scaled.floor() - (scaled.floor() - scaled).detach()
    clipped = quantized.clamp(0, self.T)
    out: torch.Tensor = clipped * self.theta / self.T
    return out

extra_repr()

Return the compact PyTorch module representation.

Source code in src/sc_neurocore/conversion/qcfs.py

def extra_repr(self) -> str:
    """Return the compact PyTorch module representation."""
    return f"T={self.T}, theta={self.theta.item():.2f}"