Merge pull request BindsNET#8 from WeihaoTan/master

Fix bugs to make CNN converted to SNN successfully

Author: djsaunde

bindsnet/conversion/conversion.py

@@ -140,6 +140,13 @@ def set_requires_grad(module, value):
                 ):
                     prev_module = module2
 
+        if isinstance(module2, nn.Linear):
+                if prev_module is not None:
+                    scale_factor = np.percentile(activations.cpu(), percentile)
+                    prev_module.weight *= prev_factor / scale_factor
+                    prev_module.bias /= scale_factor
+                    prev_factor = scale_factor
+
         else:
             activations = all_activations[name]
             if isinstance(module, nn.ReLU):
@@ -188,20 +195,16 @@ def _ann_to_snn_helper(prev, current, node_type, last=False, **kwargs):
         )
 
     elif isinstance(current, nn.Conv2d):
-        input_height, input_width = prev.shape[2], prev.shape[3]
-        out_channels, output_height, output_width = (
-            current.out_channels,
-            prev.shape[2],
-            prev.shape[3],
-        )
+        input_height, input_width = prev.shape[1], prev.shape[2]
+        out_channels = current.out_channels
 
         width = (
             input_height - current.kernel_size[0] + 2 * current.padding[0]
         ) / current.stride[0] + 1
         height = (
             input_width - current.kernel_size[1] + 2 * current.padding[1]
         ) / current.stride[1] + 1
-        shape = (1, out_channels, int(width), int(height))
+        shape = (out_channels, int(width), int(height))
 
         layer = node_type(
             shape=shape, reset=0, thresh=1, refrac=0, sum_input=last, **kwargs
@@ -223,7 +226,7 @@ def _ann_to_snn_helper(prev, current, node_type, last=False, **kwargs):
         )
 
     elif isinstance(current, nn.MaxPool2d):
-        input_height, input_width = prev.shape[2], prev.shape[3]
+        input_height, input_width = prev.shape[1], prev.shape[2]
         current.kernel_size = _pair(current.kernel_size)
         current.padding = _pair(current.padding)
         current.stride = _pair(current.stride)
@@ -234,7 +237,7 @@ def _ann_to_snn_helper(prev, current, node_type, last=False, **kwargs):
         height = (
             input_width - current.kernel_size[1] + 2 * current.padding[1]
         ) / current.stride[1] + 1
-        shape = (1, prev.shape[1], int(width), int(height))
+        shape = (prev.shape[0], int(width), int(height))
 
         layer = PassThroughNodes(shape=shape)
         connection = topology.MaxPool2dConnection(
@@ -253,7 +256,6 @@ def _ann_to_snn_helper(prev, current, node_type, last=False, **kwargs):
                 prev.shape[current.dims[0]],
                 prev.shape[current.dims[1]],
                 prev.shape[current.dims[2]],
-                prev.shape[current.dims[3]],
             ]
         )
 
@@ -265,9 +267,8 @@ def _ann_to_snn_helper(prev, current, node_type, last=False, **kwargs):
         layer = PassThroughNodes(
             shape=[
                 prev.shape[0],
-                prev.shape[1],
-                current.padding[0] + current.padding[1] + prev.shape[2],
-                current.padding[2] + current.padding[3] + prev.shape[3],
+                current.padding[0] + current.padding[1] + prev.shape[1],
+                current.padding[2] + current.padding[3] + prev.shape[2],
             ]
         )
 

bindsnet/network/network.py

@@ -392,9 +392,6 @@ def run(
                     mask=masks.get(c, None), learning=self.learning, **kwargs
                 )
 
-            # Get input to all layers.
-            current_inputs.update(self._get_inputs())
-
             # Record state variables of interest.
             for m in self.monitors:
                 self.monitors[m].record()

bindsnet/network/topology.py

@@ -285,7 +285,7 @@ def __init__(
             "(input_height - filter_height + 2 * padding_height) / stride_height + 1,"
             "(input_width - filter_width + 2 * padding_width) / stride_width + 1"
         )
-
+        
         assert (
             target.shape[0] == shape[1]
             and target.shape[1] == shape[2]
@@ -405,7 +405,7 @@ def __init__(
         self.padding = _pair(padding)
         self.dilation = _pair(dilation)
 
-        self.register_buffer("firing_rates", torch.zeros(source.shape))
+        self.register_buffer("firing_rates", torch.zeros(source.s.shape))
 
     def compute(self, s: torch.Tensor) -> torch.Tensor:
         # language=rst
@@ -429,7 +429,7 @@ def compute(self, s: torch.Tensor) -> torch.Tensor:
             return_indices=True,
         )
 
-        return s.take(indices).float()
+        return s.flatten(2).gather(2, indices.flatten(2)).view_as(indices).float()
 
     def update(self, **kwargs) -> None:
         # language=rst
@@ -452,7 +452,7 @@ def reset_state_variables(self) -> None:
         """
         super().reset_state_variables()
 
-        self.firing_rates = torch.zeros(self.source.shape)
+        self.firing_rates = torch.zeros(self.source.s.shape)
 
 
 class LocalConnection(AbstractConnection):