added some lif logic

src/neurons/leaky.rs

@@ -1,8 +1,7 @@
-#![allow(non_snake_case)]
+#![allow(non_snake_case, dead_code)]
 
 use crate::neurons::heaviside::heaviside;
 use crate::surrogate::SurrogateFn;
-use burn::module::Param;
 use burn::prelude::*;
 use burn::{Tensor, module::Module, prelude::Backend};
 
@@ -15,50 +14,71 @@ pub struct LIFConfig {
 }
 
 impl LIFConfig {
-    pub fn init<B: Backend, const D: usize>(&self, device: &B::Device) -> LIF<B, D> {
-        let initMem = Param::from_tensor(Tensor::<B, D>::zeros([1, self.neurons], device));
+    pub fn init<B: Backend, const D: usize>(&self) -> LIF {
         LIF {
             beta: self.beta,
             threshold: self.threshold,
             neurons: self.neurons,
-            hidden: initMem,
         }
     }
 }
 
 // TODO: tensor cloning and its lifecycle is probably wrong, may cause comp graph to drop. Refer burn example to find proper tensor handling..
-#[derive(Debug, Module)]
-pub struct LIF<B: Backend, const D: usize> {
+#[derive(Debug, Module, Clone)]
+pub struct LIF {
     beta: f32,
     threshold: f32,
     neurons: usize,
-    pub hidden: Param<Tensor<B, D>>,
 }
 
-impl<B: Backend, const D: usize> LIF<B, D> {
-    pub fn forward(&mut self, input: Tensor<B, D>) -> Tensor<B, D> {
-        // leaky and Integrate
-        //
-        let curMem = self.hidden.val();
-        let nxtMem = curMem.mul_scalar(self.beta).add(input.clone());
-        // fire
-        let spikes = heaviside::<B, D>(nxtMem, self.threshold, SurrogateFn::FastSigmoid);
+impl LIF {
+    pub fn forward<B: Backend, const D: usize>(
+        &mut self,
+        input: Tensor<B, D>,
+        mem: Tensor<B, D>,
+    ) -> [Tensor<B, D>; 2] {
+        // check if input shape and mem shape are same. init to zero of input shape if not.
+        if input.shape() != mem.shape() {
+            panic!(
+                "Input shape {} and memory shape {} are different",
+                input.shape(),
+                mem.shape()
+            )
+        }
 
-        self.hidden
-            .val()
-            .sub(spikes.clone().mul_scalar(self.threshold)); // requires update step fix. currently doesnt update.
-        spikes
+        // memory reset at current state.
+        let resetSignal = self.mem_reset(mem.clone());
+
+        // Decay memory and add input (B*v + X)
+        let dmem = mem.mul_scalar(self.beta).add(input);
+
+        // Reset memory based on reset method.
+        let outMem = self.step_subtract(dmem, resetSignal);
+
+        // Generate output spikes
+        let spikes = heaviside(outMem.clone(), self.threshold, SurrogateFn::FastSigmoid);
+
+        [spikes, outMem]
     }
 
-    fn mem_reset(&self, mem: Tensor<B, D>) -> Tensor<B, D> {
+    fn mem_reset<B: Backend, const D: usize>(&self, mem: Tensor<B, D>) -> Tensor<B, D> {
+        // Generates reset signal.
+        // Take diff of mem and threshold and pass through heaviside function.
         mem.sub_scalar(self.threshold)
+            .greater_elem(0.0)
+            .float()
+            .detach()
     }
 
-    pub fn reset(&mut self) {
-        // self.hidden = self.hidden.zeros_like();
+    fn step_subtract<B: Backend, const D: usize>(
+        &self,
+        input: Tensor<B, D>,
+        reset: Tensor<B, D>,
+    ) -> Tensor<B, D> {
+        input - reset.mul_scalar(self.threshold)
     }
 
-    pub fn init(&mut self, batch: usize, device: &B::Device) {
-        // self.hidden = Tensor::zeros(Shape::new([batch, self.neurons]), device);
+    pub fn reset<B: Backend, const D: usize>(&self, mem: &Tensor<B, D>) -> Tensor<B, D> {
+        Tensor::zeros_like(mem)
     }
 }