Sparse batch_eth_mnist

Author: n-shevko

bindsnet/learning/MCC_learning.py

@@ -102,7 +102,10 @@ def update(self, **kwargs) -> None:
         if ((self.min is not None) or (self.max is not None)) and not isinstance(
             self, NoOp
         ):
-            self.feature_value.clamp_(self.min, self.max)
+            if self.feature_value.is_sparse:
+                self.feature_value = self.feature_value.to_dense().clamp_(self.min, self.max).to_sparse()
+            else:
+                self.feature_value.clamp_(self.min, self.max)
 
     @abstractmethod
     def reset_state_variables(self) -> None:
@@ -247,10 +250,16 @@ def _connection_update(self, **kwargs) -> None:
                         torch.mean(self.average_buffer_pre, dim=0) * self.connection.dt
                     )
             else:
-                self.feature_value -= (
-                    self.reduction(torch.bmm(source_s, target_x), dim=0)
-                    * self.connection.dt
-                )
+                if self.feature_value.is_sparse:
+                    self.feature_value -= (
+                            torch.bmm(source_s, target_x)
+                            * self.connection.dt
+                    ).to_sparse()
+                else:
+                    self.feature_value -= (
+                        self.reduction(torch.bmm(source_s, target_x), dim=0)
+                        * self.connection.dt
+                    )
             del source_s, target_x
 
         # Post-synaptic update.
@@ -278,10 +287,16 @@ def _connection_update(self, **kwargs) -> None:
                         torch.mean(self.average_buffer_post, dim=0) * self.connection.dt
                     )
             else:
-                self.feature_value += (
-                    self.reduction(torch.bmm(source_x, target_s), dim=0)
-                    * self.connection.dt
-                )
+                if self.feature_value.is_sparse:
+                    self.feature_value += (
+                            torch.bmm(source_x, target_s)
+                            * self.connection.dt
+                    ).to_sparse()
+                else:
+                    self.feature_value += (
+                        self.reduction(torch.bmm(source_x, target_s), dim=0)
+                        * self.connection.dt
+                    )
             del source_x, target_s
 
         super().update()

bindsnet/models/models.py

@@ -4,11 +4,14 @@
 import torch
 from scipy.spatial.distance import euclidean
 from torch.nn.modules.utils import _pair
+from torch import device
 
 from bindsnet.learning import PostPre
+from bindsnet.learning.MCC_learning import PostPre as MMCPostPre
 from bindsnet.network import Network
 from bindsnet.network.nodes import DiehlAndCookNodes, Input, LIFNodes
-from bindsnet.network.topology import Connection, LocalConnection
+from bindsnet.network.topology import Connection, LocalConnection, MulticompartmentConnection
+from bindsnet.network.topology_features import Weight
 
 
 class TwoLayerNetwork(Network):
@@ -94,6 +97,9 @@ class DiehlAndCook2015(Network):
     def __init__(
         self,
         n_inpt: int,
+        device: device,
+        batch_size: int,
+        sparse: bool = False,
         n_neurons: int = 100,
         exc: float = 22.5,
         inh: float = 17.5,
@@ -169,28 +175,61 @@ def __init__(
         )
 
         # Connections
-        w = 0.3 * torch.rand(self.n_inpt, self.n_neurons)
-        input_exc_conn = Connection(
+        if sparse:
+            w = 0.3 * torch.rand(batch_size, self.n_inpt, self.n_neurons)
+        else:
+            w = 0.3 * torch.rand(self.n_inpt, self.n_neurons)
+        input_exc_conn = MulticompartmentConnection(
             source=input_layer,
             target=exc_layer,
-            w=w,
-            update_rule=PostPre,
-            nu=nu,
-            reduction=reduction,
-            wmin=wmin,
-            wmax=wmax,
-            norm=norm,
+            device=device,
+            pipeline=[
+                Weight(
+                    'weight',
+                    w,
+                    range=[wmin, wmax],
+                    norm=norm,
+                    reduction=reduction,
+                    nu=nu,
+                    learning_rule=MMCPostPre,
+                    sparse=sparse
+                )
+            ]
         )
         w = self.exc * torch.diag(torch.ones(self.n_neurons))
-        exc_inh_conn = Connection(
-            source=exc_layer, target=inh_layer, w=w, wmin=0, wmax=self.exc
+        if sparse:
+            w = w.unsqueeze(0).expand(batch_size, -1, -1)
+        exc_inh_conn = MulticompartmentConnection(
+            source=exc_layer,
+            target=inh_layer,
+            device=device,
+            pipeline=[
+                Weight(
+                    'weight',
+                    w,
+                    range=[0, self.exc],
+                    sparse=sparse
+                )
+            ]
         )
         w = -self.inh * (
             torch.ones(self.n_neurons, self.n_neurons)
             - torch.diag(torch.ones(self.n_neurons))
         )
-        inh_exc_conn = Connection(
-            source=inh_layer, target=exc_layer, w=w, wmin=-self.inh, wmax=0
+        if sparse:
+            w = w.unsqueeze(0).expand(batch_size, -1, -1)
+        inh_exc_conn = MulticompartmentConnection(
+            source=inh_layer,
+            target=exc_layer,
+            device=device,
+            pipeline=[
+                Weight(
+                    'weight',
+                    w,
+                    range=[-self.inh, 0],
+                    sparse=sparse
+                )
+            ]
         )
 
         # Add to network

bindsnet/network/topology.py

@@ -450,7 +450,11 @@ def compute(self, s: torch.Tensor) -> torch.Tensor:
             if conn_spikes.is_sparse:
                 conn_spikes = conn_spikes.to_dense()
             conn_spikes = conn_spikes.view(s.size(0), self.source.n, self.target.n)
-        out_signal = conn_spikes.sum(1)
+
+        if conn_spikes.is_sparse:
+            out_signal = conn_spikes.to_dense().sum(1)
+        else:
+            out_signal = conn_spikes.sum(1)
 
         if self.traces:
             self.activity = out_signal

examples/mnist/batch_eth_mnist.py

@@ -44,7 +44,8 @@
 parser.add_argument("--test", dest="train", action="store_false")
 parser.add_argument("--plot", dest="plot", action="store_true")
 parser.add_argument("--gpu", dest="gpu", action="store_true")
-parser.set_defaults(plot=True, gpu=True)
+parser.add_argument("--sparse", dest="sparse", action="store_true")
+parser.set_defaults(gpu=True)
 
 args = parser.parse_args()
 
@@ -66,6 +67,7 @@
 train = args.train
 plot = args.plot
 gpu = args.gpu
+sparse = args.sparse
 
 update_steps = int(n_train / batch_size / n_updates)
 update_interval = update_steps * batch_size
@@ -93,6 +95,9 @@
 
 # Build network.
 network = DiehlAndCook2015(
+    device=device,
+    sparse=sparse,
+    batch_size=batch_size,
     n_inpt=784,
     n_neurons=n_neurons,
     exc=exc,