Question about weight update

[DS-567]

I have a question regarding training using snntorch.

I am starting to work on a task where I will be needing to train a very small SNN in a rate-based task. I am starting with a simple test experiment to understand snntorch training etc using just one single LIF neuron like this:

I am attempting to pass in a frequency of anywhere between 100Hz and 1kHz as input to the neuron, and would like to train the single weight using snntorch to have the LIF neuron spike at a pre-defined frequency, like 30Hz for example. The training loop will eventually lower the weight until the exact LIF output spike count is 30 using the mse_count_loss(). I use the built-in spikegen function to produce the frequency input very easily.

Any advice on how to implement this simply using the snntorch libraries? I assume the mse_count_loss() requires labelled data?

Code so far is shown below:

########################################################################################################################################################
### LIBRARY IMPORTS ###

import snntorch as snn
import torch
from snntorch import spikeplot as splt
from snntorch import spikegen

import matplotlib.pyplot as plt

########################################################################################################################################################
### FUNCTION DECLARATIONS ###

def plot_input_output_spikes(spk_in, spk_out_rec, num_steps):

    # Generate Plots
    fig, ax = plt.subplots(2, figsize=(8,7), sharex=True, gridspec_kw = {'height_ratios': [1, 1]})

    # Plot input spikes
    splt.raster(spk_in[:,0], ax[0], s=100, c="black", marker="|")
    ax[0].set_ylabel("Input Spikes")
    ax[0].set_title("Input and Output Spikes")
    ax[0].set_yticks([])

    # Plot output spikes
    splt.raster(spk_out_rec.reshape(num_steps, -1), ax[1], s=100, c="black", marker="|")
    ax[1].set_ylabel("Output Spikes")
    ax[1].set_yticks([])

    plt.show()

########################################################################################################################################################

def get_average_frequency(spiketrain_input, spiketrain_length_input):
    spike_counter = 0
    for i in range(len(spiketrain_input)):
        if(spiketrain_input[i] == 1):
            spike_counter += 1
    average_frequency = spike_counter
    return average_frequency

########################################################################################################################################################

def input_spiketrain_generator(freq_input, spiketrain_length_input):
    rate = freq_input/1000
    spiketrain = spikegen.target_rate_code(spiketrain_length_input, rate=rate)
    return spiketrain

########################################################################################################################################################
### USER INPUT PARAMETERS ###

input_freq = 1000
output_freq = 500
num_steps = 1000
epochs = 1

########################################################################################################################################################
### VARIABLE DECLARATIONS ###

weight = torch.rand(1)

mem_lif = torch.zeros(1)
spk_lif = torch.zeros(1)

mem_lif_rec = []
spk_lif_rec = []

########################################################################################################################################################
### INPUT SPIKETRAIN GENERATION ###

input_spiketrain = input_spiketrain_generator(input_freq, num_steps)
input_spiketrain_reshaped = torch.reshape(input_spiketrain[0], (num_steps, 1))

########################################################################################################################################################
### CREATE LIF NEURON ###

lif = snn.Leaky(beta=0.8, threshold=1)

########################################################################################################################################################
### NEURON TRAINING LOOP ###

for num in range(epochs):

    for step in range(num_steps):
        spk, mem_lif = lif(input_spiketrain_reshaped[step]*weight, mem_lif)
        mem_lif_rec.append(mem_lif)
        spk_lif_rec.append(spk)

    mem_lif_rec = torch.stack(mem_lif_rec)
    spk_lif_rec = torch.stack(spk_lif_rec)

########################################################################################################################################################
### PRINT / PLOT SIMULATION INFORMATION ###

print("Weight:", weight)
print("Input average firing frequency:", get_average_frequency(input_spiketrain_reshaped, num_steps))
print("Output average firing frequency:", get_average_frequency(spk_lif_rec, num_steps))

plot_input_output_spikes(input_spiketrain_reshaped, spk_lif_rec, num_steps)