Question about STDP code in (func: stdp_conv2d_single_step)

[AndyCao1125]

Hi, I have a question when I'm runing STDP training, only using STDP for updating nn.conv2d.
However, the first error is : "UnboundLocalError: local variable 'pre_spike' referenced before assignment ".
Then, I check the code, in func: stdp_conv2d_single_step, I see

trace_pre = trace_pre - trace_pre / tau_pre + pre_spike
trace_post = trace_post - trace_post / tau_post + post_spike

But pre_spike and post_spike are defined in the below loop. I think this may cause the bug.

And I fix the code into:

  for h in range(conv.weight.shape[2]):
          for w in range(conv.weight.shape[3]):
              h_end = in_spike.shape[2] - conv.weight.shape[2] + 1 + h
              w_end = in_spike.shape[3] - conv.weight.shape[3] + 1 + w`
  
              pre_spike = in_spike[:, :, h:h_end:stride_h, w:w_end:stride_w]
              # pre_spike.shape = [N, C_in, h_out, w_out]
              post_spike = out_spike
              # post_spike.shape = [N, C_out, h_out, w_out]
              weight = conv.weight.data[:, :, h, w]
              # weight.shape = [C_out, C_in]

              ## defined here
              tr_pre = trace_pre[h][w] - trace_pre[h][w] / tau_pre + pre_spike
              tr_post = trace_post[h][w] - trace_post[h][w] / tau_post + post_spike
  
  
              # tr_pre.shape = [N, C_in, h_out, w_out]
              # tr_post.shape = [N, C_out, h_out, w_out]
              
              
              delta_w_pre = - (f_pre(weight) * 
		          (tr_post.unsqueeze(2) * pre_spike.unsqueeze(1))
                  .permute([1, 2, 0, 3, 4]).sum(dim = [2, 3, 4]))
              delta_w_post = f_post(weight) * \
                          (post_spike.unsqueeze(2) * tr_pre.unsqueeze(1))\
                          .permute([1, 2, 0, 3, 4]).sum(dim = [2, 3, 4])
              delta_w[:, :, h, w] += delta_w_pre + delta_w_post 

However, another problem is that the shape of pre_spk & post_spk always do not have the correct shape due to the post spk is sometimes generated after a maxpool layer, but this func (stdp_conv2d_single_step) can only receive the information of current nn.conv. For example, when prespk = in_spike[:, :, h:h_end:stride_h, w:w_end:stride_w] , its shape = [N, C, 32, 32] and if here we have a nn. Maxpool after the nn.conv2d, this func will receive the postspk after maxpool with shape = out_spike = [N, C,16, 16], but this function do not know. So, the calculation could not continued since the shape is not indentity.

Could you help me solve this problem? Thanks.

[fangwei123456]

@AllenYolk

[AllenYolk]

Hi AndyCao1125! I apologize for the mistakes I made in stdp_conv2d_single_step()! You may check pull request #293 for the new code.

In my original code, the following two lines are placed before the loop:

trace_pre = trace_pre - trace_pre / tau_pre + pre_spike
trace_post = trace_post - trace_post / tau_post + post_spike

However, as you said, pre_spike and post_spike are defined in the loop, and that's why UnboundLocalError: local variable 'pre_spike' referenced before assignment is raised. Actually, my intention here was to use in_spike and out_spike rather than pre_spike and post_spike, but I made a typo🤧. The modified code looks like:

if trace_pre is None:
        trace_pre = torch.zeros_like(
            in_spike, device=in_spike.device, dtype=in_spike.dtype
        )

if trace_post is None:
    trace_post = torch.zeros_like(
        out_spike, device = in_spike.device, dtype = in_spike.dtype
    )

trace_pre = trace_pre - trace_pre / tau_pre + in_spike
trace_post = trace_post - trace_post / tau_post + out_spike

The second mistake occurs in the loop. The original code for computing delta_w_post is

delta_w_post = f_post(weight) * \
                        (tr_pre.unsqueeze(1) * post_spike.unsqueeze(1))\
                        .permute([1, 2, 0, 3, 4]).sum(dim = [2, 3, 4])

and there's a shape mismatch here. The modified code looks like

delta_w_post = f_post(weight) * \
                        (tr_pre.unsqueeze(1) * post_spike.unsqueeze(2))\
                        .permute([1, 2, 0, 3, 4]).sum(dim = [2, 3, 4])

[AndyCao1125]

Thank you so much! It works in your revised version!

[AllenYolk]

I do not fully understand the "maxpooling" example you proposed🤧.

Did you mean:

1. the voltage state v of spiking neurons is first max-pooled and then the spikes are generated
2. the output of a Conv2d layer (usually considered as synaptic input to the neurons) is first max-pooled and then injected into a layer of spiking neurons
3. the output spikes of a spiking neuron layer are max-pooled

If it was case 3, you should specify the sn argument of STDPLearner as the spiking neuron layer closely following the Conv2d layer. If you meant case 1 or 2, unfortunately, the current version of STDPLearner cannot deal with these cases! STDPLearner is for updating weights in certain network architectures like:

(
    ......
    neuron.BaseNode           # generate input spikes
    nn.Linear / nn.Conv2d    # weighted connection
    neuron.BaseNode.          # generate output spikes
    ......
)

Do you have any suggestions on extending STDPLearner so that it can work in the situation you mentioned above?