Sparsity for MulticompartmentConnection

Author: n-shevko

bindsnet/network/topology.py

@@ -446,12 +446,19 @@ def compute(self, s: torch.Tensor) -> torch.Tensor:
 
         # Sum signals for each of the output/terminal neurons
         # |out_signal| = [batch_size, target.n]
-        out_signal = conn_spikes.view(s.size(0), self.source.n, self.target.n).sum(1)
+        if conn_spikes.size() != torch.Size([s.size(0), self.source.n, self.target.n]):
+            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 self.traces:
             self.activity = out_signal
 
-        return out_signal.view(s.size(0), *self.target.shape)
+        if out_signal.size() != torch.Size([s.size(0)] + self.target.shape):
+            return out_signal.view(s.size(0), *self.target.shape)
+        else:
+            return out_signal
 
     def compute_window(self, s: torch.Tensor) -> torch.Tensor:
         # language=rst

bindsnet/network/topology_features.py

@@ -31,6 +31,7 @@ def __init__(
         enforce_polarity: Optional[bool] = False,
         decay: float = 0.0,
         parent_feature=None,
+        sparse: Optional[bool] = False,
         **kwargs,
     ) -> None:
         # language=rst
@@ -47,6 +48,7 @@ def __init__(
             dimension
         :param decay: Constant multiple to decay weights by on each iteration
         :param parent_feature: Parent feature to inherit :code:`value` from
+        :param sparse: Should :code:`value` parameter be sparse tensor or not
         """
 
         #### Initialize class variables ####
@@ -61,6 +63,7 @@ def __init__(
         self.reduction = reduction
         self.decay = decay
         self.parent_feature = parent_feature
+        self.sparse = sparse
         self.kwargs = kwargs
 
         ## Backend ##
@@ -119,6 +122,10 @@ def __init__(
         self.assert_valid_range()
         if value is not None:
             self.assert_feature_in_range()
+            if self.sparse:
+                self.value = self.value.to_sparse()
+                assert not getattr(self, 'enforce_polarity', False), \
+                    "enforce_polarity isn't supported for sparse tensors"
 
     @abstractmethod
     def reset_state_variables(self) -> None:
@@ -161,7 +168,10 @@ def prime_feature(self, connection, device, **kwargs) -> None:
 
         # Check if values/norms are the correct shape
         if isinstance(self.value, torch.Tensor):
-            assert tuple(self.value.shape) == (connection.source.n, connection.target.n)
+            if self.sparse:
+                assert tuple(self.value.shape[1:]) == (connection.source.n, connection.target.n)
+            else:
+                assert tuple(self.value.shape) == (connection.source.n, connection.target.n)
 
         if self.norm is not None and isinstance(self.norm, torch.Tensor):
             assert self.norm.shape[0] == connection.target.n
@@ -214,9 +224,15 @@ def normalize(self) -> None:
         """
 
         if self.norm is not None:
-            abs_sum = self.value.sum(0).unsqueeze(0)
-            abs_sum[abs_sum == 0] = 1.0
-            self.value *= self.norm / abs_sum
+            if self.sparse:
+                abs_sum = self.value.sum(1).to_dense()
+                abs_sum[abs_sum == 0] = 1.0
+                abs_sum = abs_sum.unsqueeze(1).expand(-1, *self.value.shape[1:])
+                self.value = self.value * (self.norm / abs_sum)
+            else:
+                abs_sum = self.value.sum(0).unsqueeze(0)
+                abs_sum[abs_sum == 0] = 1.0
+                self.value *= self.norm / abs_sum
 
     def degrade(self) -> None:
         # language=rst
@@ -299,11 +315,17 @@ def assert_feature_in_range(self):
 
     def assert_valid_shape(self, source_shape, target_shape, f):
         # Multidimensional feat
-        if len(f.shape) > 1:
-            assert f.shape == (
+        if (not self.sparse and len(f.shape) > 1) or (self.sparse and len(f.shape[1:]) > 1):
+            if self.sparse:
+                f_shape = f.shape[1:]
+                expected = ('batch_size', source_shape, target_shape)
+            else:
+                f_shape = f.shape
+                expected = (source_shape, target_shape)
+            assert f_shape == (
                 source_shape,
                 target_shape,
-            ), f"Feature {self.name} has an incorrect shape of {f.shape}. Should be of shape {(source_shape, target_shape)}"
+            ), f"Feature {self.name} has an incorrect shape of {f.shape}. Should be of shape {expected}"
         # Else assume scalar, which is a valid shape
 
 
@@ -319,6 +341,7 @@ def __init__(
         reduction: Optional[callable] = None,
         decay: float = 0.0,
         parent_feature=None,
+        sparse: Optional[bool] = False
     ) -> None:
         # language=rst
         """
@@ -336,6 +359,7 @@ def __init__(
             dimension
         :param decay: Constant multiple to decay weights by on each iteration
         :param parent_feature: Parent feature to inherit :code:`value` from
+        :param sparse: Should :code:`value` parameter be sparse tensor or not
         """
 
         ### Assertions ###
@@ -349,10 +373,25 @@ def __init__(
             reduction=reduction,
             decay=decay,
             parent_feature=parent_feature,
+            sparse=sparse
+        )
+
+    def sparse_bernoulli(self):
+        values = torch.bernoulli(self.value.values())
+        mask = values != 0
+        indices = self.value.indices()[:, mask]
+        non_zero = values[mask]
+        return torch.sparse_coo_tensor(
+            indices,
+            non_zero,
+            self.value.size()
         )
 
     def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
-        return conn_spikes * torch.bernoulli(self.value)
+        if self.sparse:
+            return conn_spikes * self.sparse_bernoulli()
+        else:
+            return conn_spikes * torch.bernoulli(self.value)
 
     def reset_state_variables(self) -> None:
         pass
@@ -395,12 +434,14 @@ def __init__(
         self,
         name: str,
         value: Union[torch.Tensor, float, int] = None,
+        sparse: Optional[bool] = False
     ) -> None:
         # language=rst
         """
         Boolean mask which determines whether or not signals are allowed to traverse certain synapses.
         :param name: Name of the feature
         :param value: Boolean mask. :code:`True` means a signal can pass, :code:`False` means the synapse is impassable
+        :param sparse: Should :code:`value` parameter be sparse tensor or not
         """
 
         ### Assertions ###
@@ -419,11 +460,9 @@ def __init__(
         super().__init__(
             name=name,
             value=value,
+            sparse=sparse
         )
 
-        self.name = name
-        self.value = value
-
     def compute(self, conn_spikes) -> torch.Tensor:
         return conn_spikes * self.value
 
@@ -505,6 +544,7 @@ def __init__(
         reduction: Optional[callable] = None,
         enforce_polarity: Optional[bool] = False,
         decay: float = 0.0,
+        sparse: Optional[bool] = False
     ) -> None:
         # language=rst
         """
@@ -523,6 +563,7 @@ def __init__(
             dimension
         :param enforce_polarity: Will prevent synapses from changing signs if :code:`True`
         :param decay: Constant multiple to decay weights by on each iteration
+        :param sparse: Should :code:`value` parameter be sparse tensor or not
         """
 
         self.norm_frequency = norm_frequency
@@ -536,6 +577,7 @@ def __init__(
             nu=nu,
             reduction=reduction,
             decay=decay,
+            sparse=sparse
         )
 
     def reset_state_variables(self) -> None:
@@ -589,6 +631,7 @@ def __init__(
         value: Union[torch.Tensor, float, int] = None,
         range: Optional[Sequence[float]] = None,
         norm: Optional[Union[torch.Tensor, float, int]] = None,
+        sparse: Optional[bool] = False
     ) -> None:
         # language=rst
         """
@@ -598,13 +641,15 @@ def __init__(
         :param range: Range of acceptable values for the :code:`value` parameter
         :param norm: Value which all values in :code:`value` will sum to. Normalization of values occurs after each sample
             and after the value has been updated by the learning rule (if there is one)
+        :param sparse: Should :code:`value` parameter be sparse tensor or not
         """
 
         super().__init__(
             name=name,
             value=value,
             range=[-torch.inf, +torch.inf] if range is None else range,
             norm=norm,
+            sparse=sparse
         )
 
     def reset_state_variables(self) -> None:
@@ -629,15 +674,17 @@ def __init__(
         name: str,
         value: Union[torch.Tensor, float, int] = None,
         range: Optional[Sequence[float]] = None,
+        sparse: Optional[bool] = False
     ) -> None:
         # language=rst
         """
         Adds scalars to signals
         :param name: Name of the feature
         :param value: Values to scale signals by
+        :param sparse: Should :code:`value` parameter be sparse tensor or not
         """
 
-        super().__init__(name=name, value=value, range=range)
+        super().__init__(name=name, value=value, range=range, sparse=sparse)
 
     def reset_state_variables(self) -> None:
         pass
@@ -666,6 +713,7 @@ def __init__(
         value: Union[torch.Tensor, float, int] = None,
         degrade_function: callable = None,
         parent_feature: Optional[AbstractFeature] = None,
+        sparse: Optional[bool] = False
     ) -> None:
         # language=rst
         """
@@ -676,10 +724,11 @@ def __init__(
         :param degrade_function: Callable function which takes a single argument (:code:`value`) and returns a tensor or
         constant to be *subtracted* from the propagating spikes.
         :param parent_feature: Parent feature with desired :code:`value` to inherit
+        :param sparse: Should :code:`value` parameter be sparse tensor or not
         """
 
         # Note: parent_feature will override value. See abstract constructor
-        super().__init__(name=name, value=value, parent_feature=parent_feature)
+        super().__init__(name=name, value=value, parent_feature=parent_feature, sparse=sparse)
 
         self.degrade_function = degrade_function
 
@@ -698,6 +747,7 @@ def __init__(
         ann_values: Union[list, tuple] = None,
         const_update_rate: float = 0.1,
         const_decay: float = 0.001,
+        sparse: Optional[bool] = False
     ) -> None:
         # language=rst
         """
@@ -708,6 +758,7 @@ def __init__(
         :param value: Values to be use to build an initial mask for the synapses.
         :param const_update_rate: The mask upatate rate of the ANN decision.
         :param const_decay: The spontaneous activation of the synapses.
+        :param sparse: Should :code:`value` parameter be sparse tensor or not
         """
 
         # Define the ANN
@@ -743,16 +794,18 @@ def forward(self, x):
         self.const_update_rate = const_update_rate
         self.const_decay = const_decay
 
-        super().__init__(name=name, value=value)
+        super().__init__(name=name, value=value, sparse=sparse)
 
     def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
 
         # Update the spike buffer
         if self.start_counter == False or conn_spikes.sum() > 0:
             self.start_counter = True
-            self.spike_buffer[:, self.counter % self.spike_buffer.shape[1]] = (
-                conn_spikes.flatten()
-            )
+            if self.sparse:
+                flat_conn_spikes = conn_spikes.to_dense().flatten()
+            else:
+                flat_conn_spikes = conn_spikes.flatten()
+            self.spike_buffer[:, self.counter % self.spike_buffer.shape[1]] = flat_conn_spikes
             self.counter += 1
 
         # Update the masks
@@ -767,6 +820,8 @@ def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
 
             # self.mask = torch.clamp(self.mask, -1, 1)
             self.value = (self.mask > 0).float()
+            if self.sparse:
+                self.value = self.value.to_sparse()
 
         return conn_spikes * self.value
 
@@ -788,6 +843,7 @@ def __init__(
         ann_values: Union[list, tuple] = None,
         const_update_rate: float = 0.1,
         const_decay: float = 0.01,
+        sparse: Optional[bool] = False
     ) -> None:
         # language=rst
         """
@@ -798,6 +854,7 @@ def __init__(
         :param value: Values to be use to build an initial mask for the synapses.
         :param const_update_rate: The mask upatate rate of the ANN decision.
         :param const_decay: The spontaneous activation of the synapses.
+        :param sparse: Should :code:`value` parameter be sparse tensor or not
         """
 
         # Define the ANN
@@ -833,16 +890,18 @@ def forward(self, x):
         self.const_update_rate = const_update_rate
         self.const_decay = const_decay
 
-        super().__init__(name=name, value=value)
+        super().__init__(name=name, value=value, sparse=sparse)
 
     def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
 
         # Update the spike buffer
         if self.start_counter == False or conn_spikes.sum() > 0:
             self.start_counter = True
-            self.spike_buffer[:, self.counter % self.spike_buffer.shape[1]] = (
-                conn_spikes.flatten()
-            )
+            if self.sparse:
+                flat_conn_spikes = conn_spikes.to_dense().flatten()
+            else:
+                flat_conn_spikes = conn_spikes.flatten()
+            self.spike_buffer[:, self.counter % self.spike_buffer.shape[1]] = flat_conn_spikes
             self.counter += 1
 
         # Update the masks
@@ -857,6 +916,8 @@ def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
 
             # self.mask = torch.clamp(self.mask, -1, 1)
             self.value = (self.mask > 0).float()
+            if self.sparse:
+                self.value = self.value.to_sparse()
 
         return conn_spikes * self.value