Function dynamic_spikes

W8ASpike/quant_linear.py:11–19

(x, k=3.0)

    spike_is_available = False

def dynamic_spikes(x, k=3.0):
    vth = x.abs().mean([-1], keepdim=True).float() / k
    vth = vth.clamp(min=1e-5, max=1e4)
    spikes_int = (x / vth).round()

    if spike_is_available:
        spikes_int = spike_fake_quant(spikes_int, lif_quantizer=SpikeCountBitwiseNode(is_bidirectional=True))

    return spikes_int, vth

class QuantLinear(nn.Linear):
    def __init__(self, in_features: int, out_features: int, bias: bool = True, device=None, dtype=None, w_group_size=128, dynamic_sfr=3.0):