add qat test

modneflib/modnef/modnef_torch/model.py

@@ -92,6 +92,15 @@ class ModNEFModel(nn.Module):
 
     return super().train(mode=mode)
   
+  def init_quantizer(self):
+    """
+    initialize quantizer of laters
+    """
+
+    for m in self.modules():
+      if isinstance(m, ModNEFNeuron):
+        m.init_quantizer()
+  
   def quantize(self, force_init=False):
     """
     Quantize synaptic weight and neuron hyper-parameters
@@ -106,11 +115,11 @@ class ModNEFModel(nn.Module):
       if isinstance(m, ModNEFNeuron):
         m.quantize(force_init=force_init)
 
-  def clamp(self):
+  def clamp(self, force_init=False):
 
     for m in self.modules():
       if isinstance(m, ModNEFNeuron):
-        m.clamp()
+        m.clamp(force_init=force_init)
 
   def train(self, mode : bool = True, quant : bool = False):
     """

modneflib/modnef/modnef_torch/modnef_neurons/modnef_torch_neuron.py

@@ -150,11 +150,14 @@ class ModNEFNeuron(SpikingNeuron):
     self.quantize_weight()
     self.quantize_hp()
   
-  def clamp(self):
+  def clamp(self, force_init=False):
     """
     Clamp synaptic weight
     """
 
+    if force_init:
+      self.init_quantizer()
+
     for p in self.parameters():
       p.data = self.quantizer.clamp(p.data)
   

modneflib/modnef/modnef_torch/modnef_neurons/srlif_model/rshiftlif.py

@@ -342,7 +342,6 @@ class RShiftLIF(ModNEFNeuron):
     """
 
     self.threshold.data = self.quantizer(self.threshold.data, unscale)
-    self.beta.data = self.quantizer(self.beta.data, unscale)
 
 
   @classmethod

modneflib/modnef/modnef_torch/modnef_neurons/srlif_model/shiftlif.py

@@ -264,6 +264,9 @@ class ShiftLIF(ModNEFNeuron):
     else:
       self.mem = self.mem-self.__shift(self.mem)
 
+    if self.quantization_flag:
+      self.mem.data = self.quantizer(self.mem.data, True)
+
     if self.hardware_estimation_flag:
       self.val_min = torch.min(torch.min(input_.min(), self.mem.min()), self.val_min)
       self.val_max = torch.max(torch.max(input_.max(), self.mem.max()), self.val_max)
@@ -328,7 +331,6 @@ class ShiftLIF(ModNEFNeuron):
     """
 
     self.threshold.data = self.quantizer(self.threshold.data, unscale)
-    self.beta.data = self.quantizer(self.beta.data, unscale)
 
 
   @classmethod