modules.py

from torch import nn
import torch.nn.functional as F
import torch
from spikingjelly.clock_driven import neuron
from torch.autograd import Function

class StraightThrough(nn.Module):
    def __init__(self, channel_num: int = 1):
        super().__init__()

    def forward(self, input):
        return input


class IFNeuron(nn.Module):
    def __init__(self, scale=1.0, is_quant=False, origin_v_threshold=0, layer_type="conv", exp_device="cuda:0"):
        super(IFNeuron, self).__init__()
        self.v_threshold = nn.Parameter(scale, requires_grad=False) if type(scale) != float else scale
        self.t = 0
        self.neuron = neuron.IFNode(v_reset=None)
        if is_quant:
            self.neuron.v_threshold = scale if layer_type == "fc" else scale.unsqueeze(-1).unsqueeze(-1)
            self.neuron.v_threshold = self.neuron.v_threshold.to(exp_device)

        self.is_merged = False
        self.is_quant = is_quant
        self.origin_v_threshold = origin_v_threshold
        
    def forward(self, x):
        
        if type(self.v_threshold) == float:
            x = x / self.v_threshold
            if self.t == 0:
                self.neuron(torch.ones_like(x) * 0.5)
            x = self.neuron(x)
            self.t += 1
            if self.is_merged:
                return x
            else:
                return x * self.v_threshold 
            
        else:
    
            if self.is_quant:
                pass
            else:
                if len(x.shape) == 4: # after conv  x.shape[batch_size, channel, height, width]
                    x = x / self.v_threshold.unsqueeze(-1).unsqueeze(-1)
                else: # after Flatten  x.shape [batch_size, output]
                    x = x / self.v_threshold
            if self.t == 0:
                if self.is_quant:
                    self.neuron.v = torch.floor(torch.ones_like(x)*0.5/self.origin_v_threshold*(self.v_threshold if len(x.shape) != 4 else self.v_threshold.unsqueeze(-1).unsqueeze(-1)))
                    self.neuron.neuronal_fire()
                    self.neuron.neuronal_reset()
                else:
                    self.neuron(torch.ones_like(x)*0.5)
            
            x = self.neuron(x)
            self.t += 1
            return x 
        
    def reset(self):
        self.t = 0
        self.neuron.reset()
        

class FloorLayer(Function):
    @staticmethod
    def forward(ctx, input):
        return input.floor()

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output

qcfs = FloorLayer.apply

class QCFS(nn.Module):
    def __init__(self, up=8., t=32):
        super().__init__()
        self.up = nn.Parameter(torch.tensor([up]), requires_grad=True)
        self.t = t
    def forward(self, x):
        x = x / self.up
        x = qcfs(x*self.t+0.5)/self.t
        x = torch.clamp(x, 0, 1)
        x = x * self.up
        return x


class MPLayer(nn.Module):
    def __init__(self, v_threshold, presim_len, sim_len):
        super().__init__()        
        self.neuron = neuron.IFNode(v_reset=None)
        self.v_threshold = v_threshold
        self.t = 0
        self.membrane_lower = None
        self.presim_len = presim_len
        self.sim_len = sim_len
        
 
    def forward(self, x):
          with torch.no_grad():
              if self.t == 0:
                  self.neuron.reset()
                  self.neuron(torch.ones_like(x)*0.5)
              
              output = self.neuron(x/self.v_threshold)

              self.t += 1
              
              if self.t == self.presim_len:
                  self.membrane_lower = torch.where(self.neuron.v>1e-3,torch.ones_like(output),torch.zeros_like(output))
                  self.neuron.reset()   
                  self.neuron(torch.ones_like(x)*0.5)
              
              if self.t > self.presim_len:
                  output = output * self.membrane_lower
                                                          
              if self.t == self.presim_len + self.sim_len:                     
                  self.t = 0  
                     
              return output*self.v_threshold