SparseConnection support

Author: n-shevko

bindsnet/learning/learning.py

@@ -98,7 +98,10 @@ def update(self) -> None:
             (self.connection.wmin != -np.inf).any()
             or (self.connection.wmax != np.inf).any()
         ) and not isinstance(self, NoOp):
-            self.connection.w.clamp_(self.connection.wmin, self.connection.wmax)
+            if self.connection.w.is_sparse:
+                raise Exception("SparseConnection isn't supported for wmin\\wmax")
+            else:
+                self.connection.w.clamp_(self.connection.wmin, self.connection.wmax)
 
 
 class NoOp(LearningRule):
@@ -396,7 +399,10 @@ def _connection_update(self, **kwargs) -> None:
         if self.nu[0].any():
             source_s = self.source.s.view(batch_size, -1).unsqueeze(2).float()
             target_x = self.target.x.view(batch_size, -1).unsqueeze(1) * self.nu[0]
-            self.connection.w -= self.reduction(torch.bmm(source_s, target_x), dim=0)
+            update = self.reduction(torch.bmm(source_s, target_x), dim=0)
+            if self.connection.w.is_sparse:
+                update = update.to_sparse()
+            self.connection.w -= update
             del source_s, target_x
 
         # Post-synaptic update.
@@ -405,7 +411,10 @@ def _connection_update(self, **kwargs) -> None:
                 self.target.s.view(batch_size, -1).unsqueeze(1).float() * self.nu[1]
             )
             source_x = self.source.x.view(batch_size, -1).unsqueeze(2)
-            self.connection.w += self.reduction(torch.bmm(source_x, target_s), dim=0)
+            update = self.reduction(torch.bmm(source_x, target_s), dim=0)
+            if self.connection.w.is_sparse:
+                update = update.to_sparse()
+            self.connection.w += update
             del source_x, target_s
 
         super().update()
@@ -1113,10 +1122,14 @@ def _connection_update(self, **kwargs) -> None:
 
         # Pre-synaptic update.
         update = self.reduction(torch.bmm(source_s, target_x), dim=0)
+        if self.connection.w.is_sparse:
+            update = update.to_sparse()
         self.connection.w += self.nu[0] * update
 
         # Post-synaptic update.
         update = self.reduction(torch.bmm(source_x, target_s), dim=0)
+        if self.connection.w.is_sparse:
+            update = update.to_sparse()
         self.connection.w += self.nu[1] * update
 
         super().update()
@@ -1542,8 +1555,10 @@ def _connection_update(self, **kwargs) -> None:
             a_minus = torch.tensor(a_minus, device=self.connection.w.device)
 
         # Compute weight update based on the eligibility value of the past timestep.
-        update = reward * self.eligibility
-        self.connection.w += self.nu[0] * self.reduction(update, dim=0)
+        update = self.reduction(reward * self.eligibility, dim=0)
+        if self.connection.w.is_sparse:
+            update = update.to_sparse()
+        self.connection.w += self.nu[0] * update
 
         # Update P^+ and P^- values.
         self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus)
@@ -2214,10 +2229,11 @@ def _connection_update(self, **kwargs) -> None:
         self.eligibility_trace *= torch.exp(-self.connection.dt / self.tc_e_trace)
         self.eligibility_trace += self.eligibility / self.tc_e_trace
 
+        update = self.nu[0] * self.connection.dt * reward * self.eligibility_trace
+        if self.connection.w.is_sparse:
+            update = update.to_sparse()
         # Compute weight update.
-        self.connection.w += (
-            self.nu[0] * self.connection.dt * reward * self.eligibility_trace
-        )
+        self.connection.w += update
 
         # Update P^+ and P^- values.
         self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus)
@@ -2936,6 +2952,9 @@ def _connection_update(self, **kwargs) -> None:
         ) * source_x[:, None]
 
         # Compute weight update.
-        self.connection.w += self.nu[0] * reward * self.eligibility_trace
+        update = self.nu[0] * reward * self.eligibility_trace
+        if self.connection.w.is_sparse:
+            update = update.to_sparse()
+        self.connection.w += update
 
         super().update()

bindsnet/network/topology.py

@@ -126,9 +126,13 @@ def update(self, **kwargs) -> None:
 
         mask = kwargs.get("mask", None)
         if mask is not None:
+            if self.w.is_sparse:
+                raise Exception("Mask isn't supported for SparseConnection")
             self.w.masked_fill_(mask, 0)
 
         if self.Dales_rule is not None:
+            if self.w.is_sparse:
+                raise Exception("Dales_rule isn't supported for SparseConnection")
             # weight that are negative and should be positive are set to 0
             self.w[self.w < 0 * self.Dales_rule.to(torch.float)] = 0
             # weight that are positive and should be negative are set to 0
@@ -1947,105 +1951,12 @@ def reset_state_variables(self) -> None:
         super().reset_state_variables()
 
 
-class SparseConnection(AbstractConnection):
+class SparseConnection(Connection):
     # language=rst
     """
     Specifies sparse synapses between one or two populations of neurons.
     """
 
-    def __init__(
-        self,
-        source: Nodes,
-        target: Nodes,
-        nu: Optional[Union[float, Sequence[float], Sequence[torch.Tensor]]] = None,
-        reduction: Optional[callable] = None,
-        weight_decay: float = None,
-        **kwargs,
-    ) -> None:
-        # language=rst
-        """
-        Instantiates a :code:`Connection` object with sparse weights.
-
-        :param source: A layer of nodes from which the connection originates.
-        :param target: A layer of nodes to which the connection connects.
-         :param nu: Learning rate for both pre- and post-synaptic events. It also
-            accepts a pair of tensors to individualize learning rates of each neuron.
-            In this case, their shape should be the same size as the connection weights.
-        :param reduction: Method for reducing parameter updates along the minibatch
-            dimension.
-        :param weight_decay: Constant multiple to decay weights by on each iteration.
-
-        Keyword arguments:
-
-        :param torch.Tensor w: Strengths of synapses. Must be in ``torch.sparse`` format
-        :param float sparsity: Fraction of sparse connections to use.
-        :param LearningRule update_rule: Modifies connection parameters according to
-            some rule.
-        :param float wmin: Minimum allowed value on the connection weights.
-        :param float wmax: Maximum allowed value on the connection weights.
-        :param float norm: Total weight per target neuron normalization constant.
-        """
-        super().__init__(source, target, nu, reduction, weight_decay, **kwargs)
-
-        w = kwargs.get("w", None)
-        self.sparsity = kwargs.get("sparsity", None)
-
-        assert (
-            w is not None
-            and self.sparsity is None
-            or w is None
-            and self.sparsity is not None
-        ), 'Only one of "weights" or "sparsity" must be specified'
-
-        if w is None and self.sparsity is not None:
-            i = torch.bernoulli(
-                1 - self.sparsity * torch.ones(*source.shape, *target.shape)
-            )
-            if (self.wmin == -np.inf).any() or (self.wmax == np.inf).any():
-                v = torch.clamp(
-                    torch.rand(*source.shape, *target.shape), self.wmin, self.wmax
-                )[i.bool()]
-            else:
-                v = (
-                    self.wmin
-                    + torch.rand(*source.shape, *target.shape) * (self.wmax - self.wmin)
-                )[i.bool()]
-            w = torch.sparse.FloatTensor(i.nonzero().t(), v)
-        elif w is not None and self.sparsity is None:
-            assert w.is_sparse, "Weight matrix is not sparse (see torch.sparse module)"
-            if self.wmin != -np.inf or self.wmax != np.inf:
-                w = torch.clamp(w, self.wmin, self.wmax)
-
-        self.w = Parameter(w, requires_grad=False)
-
-    def compute(self, s: torch.Tensor) -> torch.Tensor:
-        # language=rst
-        """
-        Compute convolutional pre-activations given spikes using layer weights.
-
-        :param s: Incoming spikes.
-        :return: Incoming spikes multiplied by synaptic weights (with or without
-            decaying spike activation).
-        """
-        return torch.mm(self.w, s.view(s.shape[1], 1).float()).squeeze(-1)
-        # return torch.mm(self.w, s.unsqueeze(-1).float()).squeeze(-1)
-
-    def update(self, **kwargs) -> None:
-        # language=rst
-        """
-        Compute connection's update rule.
-        """
-
-    def normalize(self) -> None:
-        # language=rst
-        """
-        Normalize weights along the first axis according to total weight per target
-        neuron.
-        """
-
-    def reset_state_variables(self) -> None:
-        # language=rst
-        """
-        Contains resetting logic for the connection.
-        """
-        super().reset_state_variables()
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.w = Parameter(self.w.to_sparse(), requires_grad=False)