diff --git a/.gitignore b/.gitignore index c78b91a8e..c11d4ba62 100644 --- a/.gitignore +++ b/.gitignore @@ -23,5 +23,6 @@ dist/* logs/* .pytest_cache/* .vscode/* -.claude/ -data/* \ No newline at end of file +data/* +.claude/* +.claude diff --git a/bindsnet/learning/MCC_learning.py b/bindsnet/learning/MCC_learning.py index 33ebacd0a..be067e31f 100644 --- a/bindsnet/learning/MCC_learning.py +++ b/bindsnet/learning/MCC_learning.py @@ -89,8 +89,9 @@ def update(self, **kwargs) -> None: Abstract method for a learning rule update. """ - # Implement decay. - if self.decay: + # Implement decay (self.decay == 1.0 is the no-op default; skip the + # full-matrix multiply in that case). + if self.decay != 1.0: self.feature_value *= self.decay # Enforce polarities @@ -304,89 +305,92 @@ def _connection_update(self, **kwargs) -> None: def reset_state_variables(self): return - class Hebbian(MCC_LearningRule): + +class Hebbian(MCC_LearningRule): + # language=rst + """ + Simple Hebbian learning rule. Pre- and post-synaptic updates are both positive. + """ + + def __init__( + self, + connection: AbstractMulticompartmentConnection, + feature_value: Union[torch.Tensor, float, int], + range: Optional[Sequence[float]] = None, + nu: Optional[Union[float, Sequence[float]]] = None, + reduction: Optional[callable] = None, + decay: float = 0.0, + **kwargs, + ) -> None: # language=rst """ - Simple Hebbian learning rule. Pre- and post-synaptic updates are both positive. - """ + Constructor for ``Hebbian`` learning rule. - def __init__( - self, - connection: AbstractMulticompartmentConnection, - feature_value: Union[torch.Tensor, float, int], - nu: Optional[Union[float, Sequence[float]]] = None, - reduction: Optional[callable] = None, - decay: float = 0.0, + :param connection: An ``AbstractConnection`` object whose weights the + ``Hebbian`` learning rule will modify. + :param nu: Single or pair of learning rates for pre- and post-synaptic events. + :param reduction: Method for reducing parameter updates along the batch + dimension. + :param decay: Coefficient controlling rate of decay of the weights each iteration. + """ + super().__init__( + connection=connection, + feature_value=feature_value, + range=[-1, +1] if range is None else range, + nu=nu, + reduction=reduction, + decay=decay, **kwargs, - ) -> None: - # language=rst - """ - Constructor for ``Hebbian`` learning rule. - - :param connection: An ``AbstractConnection`` object whose weights the - ``Hebbian`` learning rule will modify. - :param nu: Single or pair of learning rates for pre- and post-synaptic events. - :param reduction: Method for reducing parameter updates along the batch - dimension. - :param decay: Coefficient controlling rate of decay of the weights each iteration. - """ - super().__init__( - connection=connection, - feature_value=feature_value, - nu=nu, - reduction=reduction, - decay=decay, - **kwargs, - ) + ) - assert ( - self.source.traces and self.target.traces - ), "Both pre- and post-synaptic nodes must record spike traces." + assert ( + self.source.traces and self.target.traces + ), "Both pre- and post-synaptic nodes must record spike traces." - if isinstance(MulticompartmentConnection): - self.update = self._connection_update - self.feature_value = feature_value - # elif isinstance(connection, Conv2dConnection): - # self.update = self._conv2d_connection_update - else: - raise NotImplementedError( - "This learning rule is not supported for this Connection type." - ) + if isinstance(connection, MulticompartmentConnection): + self.update = self._connection_update + self.feature_value = feature_value + # elif isinstance(connection, Conv2dConnection): + # self.update = self._conv2d_connection_update + else: + raise NotImplementedError( + "This learning rule is not supported for this Connection type." + ) - def _connection_update(self, **kwargs) -> None: - # language=rst - """ - Hebbian learning rule for ``Connection`` subclass of ``AbstractConnection`` - class. - """ + def _connection_update(self, **kwargs) -> None: + # language=rst + """ + Hebbian learning rule for ``Connection`` subclass of ``AbstractConnection`` + class. + """ - # Add polarities back to feature after updates - if self.enforce_polarity: - self.feature_value = torch.abs(self.feature_value) + # Add polarities back to feature after updates + if self.enforce_polarity: + self.feature_value = torch.abs(self.feature_value) - batch_size = self.source.batch_size + batch_size = self.source.batch_size - source_s = self.source.s.view(batch_size, -1).unsqueeze(2).float() - source_x = self.source.x.view(batch_size, -1).unsqueeze(2) - target_s = self.target.s.view(batch_size, -1).unsqueeze(1).float() - target_x = self.target.x.view(batch_size, -1).unsqueeze(1) + source_s = self.source.s.view(batch_size, -1).unsqueeze(2).float() + source_x = self.source.x.view(batch_size, -1).unsqueeze(2) + target_s = self.target.s.view(batch_size, -1).unsqueeze(1).float() + target_x = self.target.x.view(batch_size, -1).unsqueeze(1) - # Pre-synaptic update. - update = self.reduction(torch.bmm(source_s, target_x), dim=0) - self.feature_value += self.nu[0] * update + # Pre-synaptic update. + update = self.reduction(torch.bmm(source_s, target_x), dim=0) + self.feature_value += self.nu[0] * update - # Post-synaptic update. - update = self.reduction(torch.bmm(source_x, target_s), dim=0) - self.feature_value += self.nu[1] * update + # Post-synaptic update. + update = self.reduction(torch.bmm(source_x, target_s), dim=0) + self.feature_value += self.nu[1] * update - # Add polarities back to feature after updates - if self.enforce_polarity: - self.feature_value = self.feature_value * self.polarities + # Add polarities back to feature after updates + if self.enforce_polarity: + self.feature_value = self.feature_value * self.polarities - super().update() + super().update() - def reset_state_variables(self): - return + def reset_state_variables(self): + return class MSTDP(MCC_LearningRule): @@ -494,56 +498,112 @@ def _connection_update(self, **kwargs) -> None: self.target.n, device=self.target.s.device, ) - if not hasattr(self, "eligibility"): - self.eligibility = torch.zeros( - batch_size, *self.feature_value.shape, device=self.feature_value.device - ) - # Reshape pre- and post-synaptic spikes. source_s = self.source.s.view(batch_size, -1).float() target_s = self.target.s.view(batch_size, -1).float() - # Parse keyword arguments. + # Reward from current time step reward = kwargs["reward"] - a_plus = torch.tensor( - kwargs.get("a_plus", 1.0), device=self.feature_value.device + + # Build learning-rate and decay tensors, cached on (dt, device) so a + # change to either (network.dt edits, .to(device) moves) recomputes them. + dev = self.feature_value.device + dt = float(self.connection.dt) + if getattr(self, "_decay_key", None) != (dt, dev): + self._decay_key = (dt, dev) + self._a_plus_default = torch.tensor(1.0, device=dev) + self._a_minus_default = torch.tensor(-1.0, device=dev) + self._decay_plus = torch.exp(-dt / self.tc_plus.to(dev)) + self._decay_minus = torch.exp(-dt / self.tc_minus.to(dev)) + a_plus = kwargs.get("a_plus", None) + a_plus = ( + self._a_plus_default + if a_plus is None + else torch.as_tensor(a_plus, device=self.feature_value.device) + ) + a_minus = kwargs.get("a_minus", None) + a_minus = ( + self._a_minus_default + if a_minus is None + else torch.as_tensor(a_minus, device=self.feature_value.device) ) - a_minus = torch.tensor( - kwargs.get("a_minus", -1.0), device=self.feature_value.device + + # With no averaging and the standard batch reductions (squeeze for + # batch size 1, sum otherwise), the [batch, src, tgt] eligibility never + # needs to be materialized: its batch reduction is a sum of two matrix + # products applied to the weights directly. The eligibility of the past + # timestep is outer(p_plus, target_s) + outer(source_s, p_minus), with + # p_plus/p_minus still holding their previous-step values here. + fast = ( + self.average_update == 0 + and self.reduction in (torch.squeeze, torch.sum) + and not self.feature_value.is_sparse ) + if fast: + if hasattr(self, "_prev_target_s"): + if isinstance(reward, torch.Tensor): + # Keep reward on-device (no host sync for tensor rewards). + update = ( + self.p_plus.t() @ self._prev_target_s + + self._prev_source_s.t() @ self.p_minus + ) + self.feature_value += (self.nu[0].to(dev) * reward) * update + else: + alpha = float(self.nu[0]) * reward + if alpha != 0.0: + self.feature_value.addmm_( + self.p_plus.t(), self._prev_target_s, alpha=alpha + ) + self.feature_value.addmm_( + self._prev_source_s.t(), self.p_minus, alpha=alpha + ) + else: + # Dense-eligibility path: averaging buffers, custom reductions, or + # sparse weights. + if not hasattr(self, "eligibility"): + self.eligibility = torch.zeros( + batch_size, + *self.feature_value.shape, + device=self.feature_value.device, + ) - # Compute weight update based on the eligibility value of the past timestep. - update = reward * self.eligibility + # Compute weight update based on the eligibility value of the past timestep. + update = reward * self.eligibility - if self.average_update > 0: - self.average_buffer[self.average_buffer_index] = self.reduction( - update, dim=0 - ) - self.average_buffer_index = ( - self.average_buffer_index + 1 - ) % self.average_update + if self.average_update > 0: + self.average_buffer[self.average_buffer_index] = self.reduction( + update, dim=0 + ) + self.average_buffer_index = ( + self.average_buffer_index + 1 + ) % self.average_update - if self.continues_update or self.average_buffer_index == 0: - update = self.nu[0] * torch.mean(self.average_buffer, dim=0) + if self.continues_update or self.average_buffer_index == 0: + update = self.nu[0] * torch.mean(self.average_buffer, dim=0) + if self.feature_value.is_sparse: + update = update.to_sparse() + self.feature_value += update + else: + update = self.nu[0] * self.reduction(update, dim=0) if self.feature_value.is_sparse: update = update.to_sparse() self.feature_value += update - else: - update = self.nu[0] * self.reduction(update, dim=0) - if self.feature_value.is_sparse: - update = update.to_sparse() - self.feature_value += update # Update P^+ and P^- values. - self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus) + self.p_plus *= self._decay_plus self.p_plus += a_plus * source_s - self.p_minus *= torch.exp(-self.connection.dt / self.tc_minus) + self.p_minus *= self._decay_minus self.p_minus += a_minus * target_s - # Calculate point eligibility value. - self.eligibility = torch.bmm( - self.p_plus.unsqueeze(2), target_s.unsqueeze(1) - ) + torch.bmm(source_s.unsqueeze(2), self.p_minus.unsqueeze(1)) + if fast: + # Keep this step's spikes for the next step's rank-1 update. + self._prev_source_s = source_s.clone() + self._prev_target_s = target_s.clone() + else: + # Calculate point eligibility value. + self.eligibility = torch.bmm( + self.p_plus.unsqueeze(2), target_s.unsqueeze(1) + ) + torch.bmm(source_s.unsqueeze(2), self.p_minus.unsqueeze(1)) super().update() @@ -668,26 +728,36 @@ def _connection_update(self, **kwargs) -> None: # Parse keyword arguments. reward = kwargs["reward"] - a_plus = kwargs.get("a_plus", 1.0) - # if isinstance(a_plus, dict): - # for k, v in a_plus.items(): - # a_plus[k] = torch.tensor(v, device=self.feature_value.device) - # else: - a_plus = torch.tensor(a_plus, device=self.feature_value.device) - a_minus = kwargs.get("a_minus", -1.0) - # if isinstance(a_minus, dict): - # for k, v in a_minus.items(): - # a_minus[k] = torch.tensor(v, device=self.feature_value.device) - # else: - a_minus = torch.tensor(a_minus, device=self.feature_value.device) + + # Build learning-rate and decay tensors, cached on (dt, device) so a + # change to either (network.dt edits, .to(device) moves) recomputes them. + dev = self.feature_value.device + dt = float(self.connection.dt) + if getattr(self, "_decay_key", None) != (dt, dev): + self._decay_key = (dt, dev) + self._a_plus_default = torch.tensor(1.0, device=dev) + self._a_minus_default = torch.tensor(-1.0, device=dev) + self._decay_plus = torch.exp(-dt / self.tc_plus.to(dev)) + self._decay_minus = torch.exp(-dt / self.tc_minus.to(dev)) + self._decay_e_trace = torch.exp(-dt / self.tc_e_trace.to(dev)) + a_plus = kwargs.get("a_plus", None) + a_plus = ( + self._a_plus_default + if a_plus is None + else torch.as_tensor(a_plus, device=dev) + ) + a_minus = kwargs.get("a_minus", None) + a_minus = ( + self._a_minus_default + if a_minus is None + else torch.as_tensor(a_minus, device=dev) + ) # Calculate value of eligibility trace based on the value # of the point eligibility value of the past timestep. # Note: eligibility = [source.n, target.n] > 0 where source and target spiked # Note: high negs. -> - self.eligibility_trace *= torch.exp( - -self.connection.dt / self.tc_e_trace - ) # Decay + self.eligibility_trace *= self._decay_e_trace # Decay self.eligibility_trace += self.eligibility / self.tc_e_trace # Additive changes # ^ Also effected by delay in last step @@ -713,9 +783,9 @@ def _connection_update(self, **kwargs) -> None: self.feature_value += update # Update P^+ and P^- values. - self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus) # Decay + self.p_plus *= self._decay_plus # Decay self.p_plus += a_plus * source_s # Scaled source spikes - self.p_minus *= torch.exp(-self.connection.dt / self.tc_minus) # Decay + self.p_minus *= self._decay_minus # Decay self.p_minus += a_minus * target_s # Scaled target spikes # Notes: diff --git a/bindsnet/network/topology.py b/bindsnet/network/topology.py index ea3876e04..f277e831b 100644 --- a/bindsnet/network/topology.py +++ b/bindsnet/network/topology.py @@ -443,6 +443,7 @@ def __init__( pipeline: list = [], manual_update: bool = False, traces: bool = False, + sparse_compute: bool = False, **kwargs, ) -> None: # language=rst @@ -456,6 +457,11 @@ def __init__( :param manual_update: Set to :code:`True` to disable automatic updates (applying learning rules) to connection features. False by default, updates called after each time step :param traces: Set to :code:`True` to record history of connection activity (for monitors) + :param sparse_compute: Set to :code:`True` to read only the rows of the effective + weight for source neurons that spiked. A win when few sources are active; + on CUDA it is applied only for large connections + (``source.n * target.n >= 4e6``), where the required device sync pays + for itself. Ignored otherwise. """ super().__init__(source, target, device, pipeline, **kwargs) @@ -464,49 +470,114 @@ def __init__( if self.traces: self.activity = None + self.sparse_compute = sparse_compute + + # Cached (a_eff, b_sum) for pipelines whose features are all static + # (see AbstractFeature.is_static). Invalidated whenever a feature can + # change: learning updates, normalize, reset, device/dtype moves. + self._fold_cache = None + + def _apply(self, fn, recurse=True): + self._fold_cache = None + return super()._apply(fn, recurse) + def compute(self, s: torch.Tensor) -> torch.Tensor: # language=rst """ - Compute pre-activations given spikes using connection weights. - - :param s: Incoming spikes. - :return: Incoming spikes multiplied by synaptic weights (with or without - decaying spike activation). - """ + Direct incoming spikes through the connection's feature pipeline. - # Change to numeric type (torch doesn't like booleans for matrix ops) - # Note: .float() is an expensive operation. Use as minimally as possible! - # if s.dtype != torch.float32: - # s = s.float() + Each feature's ``compute`` returns its ``[source.n, target.n]`` value; how + it folds is set by the feature's ``op`` (``"mul"`` default, ``"add"``, + ``"sub"``). Folding the recurrence (start ``A = 1``, ``B = 0``): - # Prepare broadcast from incoming spikes to all output neurons - # |conn_spikes| = [batch_size, source.n * target.n] - conn_spikes = s.view(s.size(0), self.source.n, 1).repeat(1, 1, self.target.n) - # TODO: ^ This could probably be optimized + * ``mul`` factor ``a``: ``A <- a * A`` and ``B <- a * B`` + * ``add`` term ``b``: ``B <- B + b`` + * ``sub`` term ``b``: ``B <- B - b`` - # Run through pipeline - for f in self.pipeline: - conn_spikes = f.compute(conn_spikes) + ``B`` stays ``None`` (unallocated) unless an additive feature is present, + so a purely multiplicative pipeline is exactly the single ``s @ A`` matmul. - # Sum signals for each of the output/terminal neurons - # |out_signal| = [batch_size, target.n] - 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) + :param s: Incoming spikes, shape ``[batch, *source.shape]``. + :return: Post-synaptic input, shape ``[batch, *target.shape]``. + """ + s = s.view(s.size(0), self.source.n) - if conn_spikes.is_sparse: - out_signal = conn_spikes.to_dense().sum(1) + deferred = None + if self._fold_cache is not None: + a_eff, b_sum = self._fold_cache + else: + # running product of multiplicative factors, [source.n, target.n] + a_eff = None + # running additive offset, [source.n, target.n]; None while still zero + b_eff = None + for f in self.pipeline: + factor = f.compute(s) + # Compute-time side effects (e.g. per-time-step weight + # normalization) run after the fold has consumed this value. + d = getattr(f, "defer", None) + if d is not None: + deferred = [d] if deferred is None else deferred + [d] + if factor is None: + # Side-effect-only pipeline entries (sub-features) fold as identity. + continue + if isinstance(factor, torch.Tensor) and factor.is_sparse: + # Sparse feature values carry a leading batch dim ([1, src, tgt]) + # from prime_feature; densify to the fold's [src, tgt] shape. + factor = factor.to_dense().view(self.source.n, self.target.n) + op = getattr(f, "op", "mul") + if op == "mul": + a_eff = factor if a_eff is None else a_eff * factor + if b_eff is not None: + b_eff = b_eff * factor + else: # additive contribution: "add" -> +factor, "sub" -> -factor + term = factor if op == "add" else -factor + b_eff = term if b_eff is None else b_eff + term + + if a_eff is None: + # Degenerate pipeline with no multiplicative feature: every source + # neuron contributes with unit weight. + a_eff = torch.ones(self.source.n, self.target.n, device=s.device) + if not torch.is_floating_point(a_eff): + a_eff = a_eff.float() + + # Additive terms apply to every synapse regardless of spikes, so + # their contribution is the source-sum, a constant [target.n] row. + b_sum = b_eff.sum(dim=0) if b_eff is not None else None + + if all(getattr(f, "is_static", True) for f in self.pipeline): + self._fold_cache = (a_eff, b_sum) + + # The gather pays for its ``nonzero()`` only where that sync is expensive + # relative to the matmul: on CUDA it needs a large weight matrix to win; + # on CPU there is no device sync, so it helps even at small sizes. + use_gather = self.sparse_compute and ( + not a_eff.is_cuda or self.source.n * self.target.n >= 4_000_000 + ) + if use_gather: + # Read only the rows of A for source neurons that spiked this step + # (numerically identical; faster only when few are active). + active = s.any(dim=0).nonzero(as_tuple=False).squeeze(1) + if active.numel() == 0: + out = torch.zeros( + s.size(0), a_eff.size(-1), device=a_eff.device, dtype=a_eff.dtype + ) + else: + out = s[:, active].to(a_eff.dtype) @ a_eff.index_select(0, active) else: - out_signal = conn_spikes.sum(1) + out = s.to(a_eff.dtype) @ a_eff - if self.traces: - self.activity = out_signal + if b_sum is not None: + out = out + b_sum - 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 + if deferred is not None: + for fn in deferred: + fn() + + if self.traces: + self.activity = out + if out.size() != torch.Size([s.size(0)] + self.target.shape): + return out.view(s.size(0), *self.target.shape) + return out def compute_window(self, s: torch.Tensor) -> torch.Tensor: # language=rst @@ -544,6 +615,7 @@ def update(self, **kwargs) -> None: learning = kwargs.get("learning", False) if learning and not self.manual_update: # Pipeline learning + self._fold_cache = None for f in self.pipeline: f.update(**kwargs) @@ -553,6 +625,7 @@ def normalize(self) -> None: Normalize all features in the connection. """ # Normalize pipeline features + self._fold_cache = None for f in self.pipeline: f.normalize() @@ -563,6 +636,7 @@ def reset_state_variables(self) -> None: """ super().reset_state_variables() + self._fold_cache = None for f in self.pipeline: f.reset_state_variables() diff --git a/bindsnet/network/topology_features.py b/bindsnet/network/topology_features.py index 43ccc1389..7abcd6368 100644 --- a/bindsnet/network/topology_features.py +++ b/bindsnet/network/topology_features.py @@ -18,6 +18,19 @@ class AbstractFeature(ABC): Features to operate on signals traversing a connection. """ + # How this feature folds in the connection's affine pipeline (see + # :meth:`MulticompartmentConnection.compute`): ``"mul"`` (elementwise multiply, + # the default), ``"add"`` or ``"sub"`` (elementwise add/subtract of the value + # returned by :meth:`compute`). + op = "mul" + + # Whether :meth:`compute` returns the same value every step until the + # feature is mutated through the connection (learning update, normalize, + # reset, device move). Static-only pipelines let the connection cache the + # folded factors between mutations. Features whose value depends on the + # incoming spikes or is resampled per step must set this to ``False``. + is_static = True + @abstractmethod def __init__( self, @@ -78,6 +91,7 @@ def __init__( from ..learning.MCC_learning import ( NoOp, PostPre, + Hebbian, MSTDP, MSTDPET, ) @@ -85,6 +99,7 @@ def __init__( supported_rules = [ NoOp, PostPre, + Hebbian, MSTDP, MSTDPET, ] @@ -114,9 +129,9 @@ def __init__( ), "Feature {0}'s nu should be of type list or tuple, not {1}".format( name, type(nu) ) - assert reduction is None or isinstance( - reduction, callable - ), "Feature {0}'s reduction should be of type callable, not {1}".format( + assert reduction is None or callable( + reduction + ), "Feature {0}'s reduction should be callable, not {1}".format( name, type(reduction) ) assert decay is None or isinstance( @@ -163,10 +178,13 @@ def reset_state_variables(self) -> None: pass @abstractmethod - def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]: + def compute(self, s) -> Union[torch.Tensor, float, int]: # language=rst """ - Computes the feature being operated on a set of incoming signals. + Return this feature's ``[source.n, target.n]`` value, given pre-synaptic + spikes ``s`` of shape ``[batch, source.n]``. How the value folds into the + connection is set by :attr:`op` -- a multiplicative factor (``"mul"``, the + default), or an additive/subtractive offset (``"add"``/``"sub"``). """ pass @@ -422,11 +440,14 @@ def sparse_bernoulli(self): 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]: + # Resampled every step; never cacheable. + is_static = False + + def compute(self, s) -> Union[torch.Tensor, float, int]: + # Factor: a fresh Bernoulli draw each step (resampled, never cached). if self.sparse: - return conn_spikes * self.sparse_bernoulli() - else: - return conn_spikes * torch.bernoulli(self.value) + return self.sparse_bernoulli() + return torch.bernoulli(self.value) def reset_state_variables(self) -> None: pass @@ -504,8 +525,8 @@ def __init__( self.name = name self.value = value - def compute(self, conn_spikes) -> torch.Tensor: - return conn_spikes * self.value + def compute(self, s) -> torch.Tensor: + return self.value def reset_state_variables(self) -> None: pass @@ -550,6 +571,9 @@ def prime_feature(self, connection, device, **kwargs) -> None: class MeanField(AbstractFeature): + # Depends on the incoming spikes; never cacheable. + is_static = False + def __init__(self) -> None: # language=rst """ @@ -560,9 +584,9 @@ def __init__(self) -> None: def reset_state_variables(self) -> None: pass - def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]: - return conn_spikes.mean() * torch.ones( - self.source_n * self.target_n, device=self.device + def compute(self, s) -> Union[torch.Tensor, float, int]: + return s.float().mean() * torch.ones( + self.source_n, self.target_n, device=s.device ) def prime_feature(self, connection, device, **kwargs) -> None: @@ -613,6 +637,12 @@ def __init__( self.norm_frequency = norm_frequency self.enforce_polarity = enforce_polarity + if norm_frequency == "time step": + # Normalization mutates ``value`` on every compute; the connection + # runs ``defer`` after the fold has consumed the pre-normalization + # value (this avoids cloning the full matrix every step). + self.is_static = False + self.defer = lambda: self.normalize(time_step_norm=True) super().__init__( name=name, value=value, @@ -630,7 +660,7 @@ def __init__( def reset_state_variables(self) -> None: pass - def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]: + def compute(self, s) -> Union[torch.Tensor, float, int]: if self.enforce_polarity: pos_mask = ~torch.logical_xor(self.value > 0, self.positive_mask) neg_mask = ~torch.logical_xor(self.value < 0, ~self.positive_mask) @@ -638,11 +668,7 @@ def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]: self.value[~pos_mask] = 0.0001 self.value[~neg_mask] = -0.0001 - return_val = self.value * conn_spikes - if self.norm_frequency == "time step": - self.normalize(time_step_norm=True) - - return return_val + return self.value def prime_feature(self, connection, device, **kwargs) -> None: #### Initialize value #### @@ -705,11 +731,15 @@ def __init__( batch_size=batch_size, ) + # Bias is additive: folds as ``B <- B + value`` in the connection's pipeline. + op = "add" + def reset_state_variables(self) -> None: pass - def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]: - return conn_spikes + self.value + def compute(self, s) -> Union[torch.Tensor, float, int]: + # Additive offset added to every synapse (independent of the spikes). + return self.value def prime_feature(self, connection, device, **kwargs) -> None: #### Initialize value #### @@ -733,7 +763,7 @@ def __init__( ) -> None: # language=rst """ - Adds scalars to signals + Multiply all signals by a scalar :param name: Name of the feature :param value: Values to scale signals by :param value_dtype: Data type for :code:`value` tensor @@ -752,8 +782,8 @@ def __init__( def reset_state_variables(self) -> None: pass - def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]: - return conn_spikes * self.value + def compute(self, s) -> Union[torch.Tensor, float, int]: + return self.value def prime_feature(self, connection, device, **kwargs) -> None: #### Initialize value #### @@ -806,14 +836,23 @@ def __init__( self.degrade_function = degrade_function + # Degradation is subtractive: folded as ``B <- B - degrade_function(value)``. + op = "sub" + def reset_state_variables(self) -> None: pass - def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]: - return conn_spikes - self.degrade_function(self.value) + def compute(self, s) -> Union[torch.Tensor, float, int]: + # Subtractive offset (via degrade_function) applied to every synapse. + if self.degrade_function is not None: + return self.degrade_function(self.value) + return self.value class AdaptationBaseSynapsHistory(AbstractFeature): + # Value evolves with the spike history every step; never cacheable. + is_static = False + def __init__( self, name: str, @@ -882,7 +921,11 @@ def forward(self, x): batch_size=batch_size, ) - def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]: + def compute(self, s) -> Union[torch.Tensor, float, int]: + # This feature needs the per-synapse spikes, so build them from s + conn_spikes = s.view(s.size(0), -1, 1).expand( + s.size(0), s.size(1), self.value.size(-1) + ) # Update the spike buffer if self.start_counter == False or conn_spikes.sum() > 0: @@ -911,7 +954,7 @@ def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]: if self.sparse: self.value = self.value.to_sparse() - return conn_spikes * self.value + return self.value def reset_state_variables( self, @@ -924,6 +967,9 @@ def reset_state_variables( class AdaptationBaseOtherSynaps(AbstractFeature): + # Value evolves with the spike history every step; never cacheable. + is_static = False + def __init__( self, name: str, @@ -992,7 +1038,11 @@ def forward(self, x): batch_size=batch_size, ) - def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]: + def compute(self, s) -> Union[torch.Tensor, float, int]: + # This feature needs the per-synapse spikes, so build them from s + conn_spikes = s.view(s.size(0), -1, 1).expand( + s.size(0), s.size(1), self.value.size(-1) + ) # Update the spike buffer if self.start_counter == False or conn_spikes.sum() > 0: @@ -1021,7 +1071,7 @@ def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]: if self.sparse: self.value = self.value.to_sparse() - return conn_spikes * self.value + return self.value def reset_state_variables( self, @@ -1043,6 +1093,10 @@ class AbstractSubFeature(ABC): execution. """ + # Runs a side effect on every step; the pipeline must never be cached + # around it. + is_static = False + @abstractmethod def __init__( self, @@ -1060,15 +1114,16 @@ def __init__( self.parent = parent_feature self.sub_feature = None # <-- Defined in non-abstract constructor - def compute(self, _) -> None: + def compute(self, s): # language=rst """ - Proxy function to catch a pipeline execution from topology.py's :code:`compute` function. Allows :code:`SubFeature` - objects to be executed like real features in the pipeline. + Proxy to run a parent feature's side-effect (e.g. normalize/update) from + inside the pipeline. Returns ``None`` so the fold skips it entirely. """ # sub_feature should be defined in the non-abstract constructor self.sub_feature() + return None class Normalization(AbstractSubFeature): diff --git a/examples/benchmark/sparse_compute and foldable pipeline/_bench_common.py b/examples/benchmark/sparse_compute and foldable pipeline/_bench_common.py new file mode 100644 index 000000000..91c21cffe --- /dev/null +++ b/examples/benchmark/sparse_compute and foldable pipeline/_bench_common.py @@ -0,0 +1,142 @@ +""" +Shared machinery for the ExampleNetwork MCC benchmarks (imported by +``foldable_pipelines.py``, ``sparse_compute.py`` and ``both.py``). + +Each benchmark reports the %-time speedup of one MCC configuration over another by +timing the *same* ExampleNetwork under different ``compute`` modes: + + * ``expansion`` -- the pre-optimization path, reconstructed here as a + monkeypatch: materialize ``[batch, src, tgt]``, apply every + feature elementwise, then sum over source. (This path was + removed from the code when the fold landed, so we rebuild it + to serve as the baseline.) + * ``fold`` -- the current folded path: ``out = s @ A + B.sum(0)``. + * ``fold_sparse`` -- the fold plus activity-sparse compute (``sparse_compute``): + read only the weight rows of source neurons that spiked. + +Speedup is wall-time reduction: ``(t_baseline - t_new) / t_baseline * 100`` +(positive = faster). Learning is disabled so we time the pure forward path (the +part these optimizations affect). +""" + +import os +import statistics +import sys +import time + +_HERE = os.path.dirname(os.path.abspath(__file__)) # .../examples/benchmark/ +_EXAMPLES = os.path.dirname(os.path.dirname(_HERE)) # .../examples +_ROOT = os.path.dirname(_EXAMPLES) # repo root (for ``bindsnet``) +_STRESS = os.path.join(_EXAMPLES, "stress_test") # for ``example_network`` +for _p in (_ROOT, _STRESS): + if _p not in sys.path: + sys.path.insert(0, _p) + +import torch + +from bindsnet.network.topology import MulticompartmentConnection +from bindsnet.network.topology_features import Degradation, Probability +from example_network import ExampleNetwork + +# ExampleNetwork sizes per device: 20k excitatory neurons on GPU (where the fold +# shines), a smaller net on CPU so the baseline finishes in reasonable time. +GPU_CONFIG = dict(in_size=100, exc_size=20_000, inh_size=2_000) +CPU_CONFIG = dict(in_size=100, exc_size=2_000, inh_size=200) +GPU_TIME, CPU_TIME = 50, 20 +REPS, WARMUP = 5, 2 + + +def _expansion_compute(self, s): + """The pre-fold ``[batch, src, tgt]`` expansion path (baseline).""" + s = s.view(s.size(0), self.source.n) + cs = s.view(s.size(0), self.source.n, 1).repeat(1, 1, self.target.n) + for f in self.pipeline: + op = getattr(f, "op", "mul") + if isinstance(f, Probability): + v = torch.bernoulli(f.value) + elif isinstance(f, Degradation): + v = ( + f.degrade_function(f.value) + if f.degrade_function is not None + else f.value + ) + else: + v = f.value + if op == "mul": + cs = cs * v + elif op == "add": + cs = cs + v + else: + cs = cs - v + out = cs.sum(1) + if getattr(self, "traces", False): + self.activity = out + if out.size() != torch.Size([s.size(0)] + self.target.shape): + return out.view(s.size(0), *self.target.shape) + return out + + +def _set_mode(net, mode): + for c in net.connections.values(): + if isinstance(c, MulticompartmentConnection): + if mode == "expansion": + c.compute = _expansion_compute.__get__(c) # instance override + c.sparse_compute = False + else: + c.__dict__.pop("compute", None) # restore the class (fold) method + c.sparse_compute = mode == "fold_sparse" + + +def _bench(net, inputs, T, device, modes): + cuda = device.startswith("cuda") + net.train(False) # forward-only: time the compute path the optimizations touch + for m in modes: # warmup each mode + _set_mode(net, m) + for _ in range(WARMUP): + net.reset_state_variables() + net.run(inputs=inputs, time=T) + if cuda: + torch.cuda.synchronize() + samples = {m: [] for m in modes} + for _ in range(REPS): # interleave modes each rep to counter thermal drift + for m in modes: + _set_mode(net, m) + net.reset_state_variables() + if cuda: + torch.cuda.synchronize() + t0 = time.perf_counter() + net.run(inputs=inputs, time=T) + if cuda: + torch.cuda.synchronize() + samples[m].append(time.perf_counter() - t0) + return {m: statistics.median(v) * 1e3 for m, v in samples.items()} + + +def run_speedup(baseline_mode, test_mode, technique): + """Build the ExampleNetwork on CPU and GPU, time both modes, print the speedup.""" + print("=" * 72) + print(f"{technique}") + print(f" (%-time speedup of '{test_mode}' vs baseline '{baseline_mode}')") + print("=" * 72) + + devices = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + if not torch.cuda.is_available(): + print("[note] CUDA not available -> reporting CPU only.\n") + + for dev in devices: + cfg = GPU_CONFIG if dev == "cuda" else CPU_CONFIG + T = GPU_TIME if dev == "cuda" else CPU_TIME + net = ExampleNetwork(device=dev, **cfg) + inputs = net.make_input(T) + res = _bench(net, inputs, T, dev, [baseline_mode, test_mode]) + base, new = res[baseline_mode], res[test_mode] + speedup = (base - new) / base * 100.0 + print( + f" [{dev.upper():4s}] exc={cfg['exc_size']:>6d} time={T:>3d} | " + f"{baseline_mode}={base:8.2f} ms {test_mode}={new:8.2f} ms " + f"| speedup = {speedup:+6.1f}% time" + ) + del net + if torch.cuda.is_available(): + torch.cuda.empty_cache() + print() diff --git a/examples/benchmark/sparse_compute and foldable pipeline/foldable_pipelines.py b/examples/benchmark/sparse_compute and foldable pipeline/foldable_pipelines.py new file mode 100644 index 000000000..8295abaa2 --- /dev/null +++ b/examples/benchmark/sparse_compute and foldable pipeline/foldable_pipelines.py @@ -0,0 +1,9 @@ +from _bench_common import run_speedup + +### Benchmark: foldable pipelines ### +# Measures the speedup of the folded MultiCompartmentConnection compute (``out = s @ A + B.sum(0)``) over +# the pre-fold ``[batch, src, tgt]`` expansion path, on the ExampleNetwork +if __name__ == "__main__": + run_speedup( + baseline_mode="expansion", test_mode="fold", technique="Foldable pipelines" + ) diff --git a/examples/benchmark/sparse_compute and foldable pipeline/run_both.py b/examples/benchmark/sparse_compute and foldable pipeline/run_both.py new file mode 100644 index 000000000..0752da254 --- /dev/null +++ b/examples/benchmark/sparse_compute and foldable pipeline/run_both.py @@ -0,0 +1,9 @@ +from _bench_common import run_speedup + +# Run both sparse_compute.py and foldable_pipelines.py +if __name__ == "__main__": + run_speedup( + baseline_mode="expansion", + test_mode="fold_sparse", + technique="Foldable pipelines + sparse compute (combined)", + ) diff --git a/examples/benchmark/sparse_compute and foldable pipeline/sparse_compute.py b/examples/benchmark/sparse_compute and foldable pipeline/sparse_compute.py new file mode 100644 index 000000000..5781e8918 --- /dev/null +++ b/examples/benchmark/sparse_compute and foldable pipeline/sparse_compute.py @@ -0,0 +1,10 @@ +from _bench_common import run_speedup + +### Benchmark: sparse compute ### +# Measures the speedup of activity-sparse compute (``sparse_compute=True`` -- read +# only the weight rows of source neurons that spiked) over the dense folded compute, +# on the ExampleNetwork (CPU and GPU). +if __name__ == "__main__": + run_speedup( + baseline_mode="fold", test_mode="fold_sparse", technique="Sparse compute" + ) diff --git a/examples/stress_test/example_network.py b/examples/stress_test/example_network.py new file mode 100644 index 000000000..3d8e94f97 --- /dev/null +++ b/examples/stress_test/example_network.py @@ -0,0 +1,183 @@ +""" +ExampleNetwork -- a large, sparse, recurrent ``MulticompartmentConnection`` (MCC) +stress workload. + +Topology: ``Input(I) -> EXC_LIF <-> INH_LIF``. Four single-``Weight`` MCC. +The ``I -> EXC`` connection also carries an ``MSTDP`` learning rule. + +Run it directly to stress the simulator: + + python examples/stress_test/example_network.py --device cuda --exc 20000 --time 50 + python examples/stress_test/example_network.py --device cpu --exc 2000 --time 20 +""" + +import argparse +import os +import sys +import time + +# Make ``bindsnet`` importable when this file is run as a standalone script. +sys.path.insert( + 0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +) + +import torch + +from bindsnet.learning.MCC_learning import MSTDP +from bindsnet.network.network import Network +from bindsnet.network.nodes import Input, LIFNodes +from bindsnet.network.topology import MulticompartmentConnection +from bindsnet.network.topology_features import Weight + + +class ExampleNetwork(Network): + def __init__( + self, + device="cpu", + in_size=100, + exc_size=20_000, + inh_size=2_000, + batch_size=1, + i_to_exc_connectivity=0.15, + i_to_inh_connectivity=0.05, + inh_to_exc_connectivity=0.05, + exc_to_inh_connectivity=0.05, + ): + super().__init__() + self.device = device + self.in_size = in_size + self.exc_size = exc_size + self.inh_size = inh_size + self.batch_size = batch_size + self.i_to_exc_connectivity = i_to_exc_connectivity + self.i_to_inh_connectivity = i_to_inh_connectivity + self.inh_to_exc_connectivity = inh_to_exc_connectivity + self.exc_to_inh_connectivity = exc_to_inh_connectivity + self.build() + + def _sparse_weight(self, rows, cols, connectivity, sign=1.0): + w = sign * torch.rand(rows, cols, device=self.device) + keep = torch.rand(rows, cols, device=self.device) > (1 - connectivity) + return w * keep + + def build(self): + device = self.device + self.add_layer(layer=Input(self.in_size), name="I") + self.add_layer(layer=LIFNodes(self.exc_size), name="EXC_LIF") + self.add_layer(layer=LIFNodes(self.inh_size), name="INH_LIF") + self.add_connection( + connection=MulticompartmentConnection( + source=self.layers["I"], + target=self.layers["EXC_LIF"], + device=device, + pipeline=[ + Weight( + name="I_to_EXC_weight", + value=self._sparse_weight( + self.in_size, self.exc_size, self.i_to_exc_connectivity + ), + learning_rule=MSTDP, + range=(0, 1), + ) + ], + ), + source="I", + target="EXC_LIF", + ) + self.add_connection( + connection=MulticompartmentConnection( + source=self.layers["I"], + target=self.layers["INH_LIF"], + device=device, + pipeline=[ + Weight( + name="I_to_INH_weight", + value=self._sparse_weight( + self.in_size, self.inh_size, self.i_to_inh_connectivity + ), + ) + ], + ), + source="I", + target="INH_LIF", + ) + self.add_connection( + connection=MulticompartmentConnection( + source=self.layers["INH_LIF"], + target=self.layers["EXC_LIF"], + device=device, + pipeline=[ + Weight( + name="INH_to_EXC_weight", + value=self._sparse_weight( + self.inh_size, + self.exc_size, + self.inh_to_exc_connectivity, + sign=-1.0, + ), + ) + ], + ), + source="INH_LIF", + target="EXC_LIF", + ) + self.add_connection( + connection=MulticompartmentConnection( + source=self.layers["EXC_LIF"], + target=self.layers["INH_LIF"], + device=device, + pipeline=[ + Weight( + name="EXC_to_INH_weight", + value=self._sparse_weight( + self.exc_size, self.inh_size, self.exc_to_inh_connectivity + ), + ) + ], + ), + source="EXC_LIF", + target="INH_LIF", + ) + self.to(device) + + def make_input(self, runtime): + # Poisson-ish random spike train into the input layer. + return { + "I": torch.rand(runtime, self.batch_size, self.in_size, device=self.device) + > 0.90 + } + + +if __name__ == "__main__": + p = argparse.ArgumentParser(description="Stress-test the ExampleNetwork.") + p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + p.add_argument("--in-size", type=int, default=100) + p.add_argument("--exc", type=int, default=20_000) + p.add_argument("--inh", type=int, default=2_000) + p.add_argument("--time", type=int, default=50) + args = p.parse_args() + + device = args.device + if device.startswith("cuda") and not torch.cuda.is_available(): + print("[note] CUDA unavailable; falling back to CPU.") + device = "cpu" + + net = ExampleNetwork( + device=device, in_size=args.in_size, exc_size=args.exc, inh_size=args.inh + ) + net.train(False) # forward-only stress (no learning) + inputs = net.make_input(args.time) + + net.run(inputs=inputs, time=args.time) # warmup + if device.startswith("cuda"): + torch.cuda.synchronize() + net.reset_state_variables() + t0 = time.perf_counter() + net.run(inputs=inputs, time=args.time) + if device.startswith("cuda"): + torch.cuda.synchronize() + ms = (time.perf_counter() - t0) * 1e3 + print( + f"ExampleNetwork [{device}] in={args.in_size} exc={args.exc} inh={args.inh} " + f"time={args.time}: {ms:.1f} ms total ({ms / args.time:.3f} ms/step)" + ) diff --git a/test/network/test_connections.py b/test/network/test_connections.py index 24f2333e8..fec421da2 100644 --- a/test/network/test_connections.py +++ b/test/network/test_connections.py @@ -1,4 +1,5 @@ import torch +import math from bindsnet.learning import ( MSTDP, @@ -12,6 +13,8 @@ from bindsnet.network import Network from bindsnet.network.nodes import Input, LIFNodes, SRM0Nodes from bindsnet.network.topology import * +import bindsnet.learning.MCC_learning as mcc +import bindsnet.network.topology_features as tf class TestConnection: @@ -19,12 +22,7 @@ class TestConnection: Tests all stable groups of neurons / nodes. """ - def __init__(self): - if torch.cuda.is_available(): - self.device = torch.device("cuda:0") - else: - self.device = torch.device("cpu:0") - print(f"Using device '{self.device}' for the test") + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") def test_transfer(self): if not torch.cuda.is_available(): @@ -45,7 +43,7 @@ def test_transfer(self): l_b = LIFNodes(shape=[1, 26, 26]) connection = conn_type(l_a, l_b, *args, **kwargs) - connection.to() + connection.to(self.device) connection_tensors = [ k @@ -70,7 +68,9 @@ def test_transfer(self): print(d, d == torch.device("cuda:0")) assert d == torch.device("cuda:0") - def test_weights(self, conn_type, shape_a, shape_b, shape_w, *args, **kwargs): + # Not named test_*: this is a manual matrix check driven from __main__ (it + # takes arguments, so pytest cannot collect it). + def check_weights(self, conn_type, shape_a, shape_b, shape_w, *args, **kwargs): print("Testing:", conn_type) time = 100 weights = [None, torch.Tensor(*shape_w)] @@ -164,7 +164,622 @@ def test_weights(self, conn_type, shape_a, shape_b, shape_w, *args, **kwargs): ) +class TestMultiCompartmentConnection: + + device = torch.device("cpu") + + # ----------------------------------------------------------------------- # + # Helpers # + # ----------------------------------------------------------------------- # + + def _make_mcc(self, pipeline, src_n, tgt_n, batch=1, sparse_compute=False): + """Build (and prime) a standalone MCC: Input(src_n) -> LIFNodes(tgt_n).""" + src = Input(n=src_n, traces=True) + tgt = LIFNodes(n=tgt_n, traces=True) + # batch_size is None until a Network sets it; the learning rules need it. + src.batch_size = batch + tgt.batch_size = batch + conn = MulticompartmentConnection( + source=src, + target=tgt, + device=self.device, + pipeline=pipeline, + sparse_compute=sparse_compute, + ) + conn.dt = 1.0 # not set until added to a Network; rules read connection.dt + return conn + + def _reference_expansion(self, pipeline, s, tgt_n): + """Pre-collapse pipeline features""" + b, src = s.shape + x = s.view(b, src, 1).expand(b, src, tgt_n).clone().float() + for f in pipeline: + op = getattr(f, "op", "mul") + if isinstance(f, tf.Degradation): + v = ( + f.degrade_function(f.value) + if f.degrade_function is not None + else f.value + ) + else: + v = f.value + if torch.is_tensor(v): + v = v.float() + if op == "mul": + x = x * v + elif op == "add": + x = x + v + else: # "sub" + x = x - v + return x.sum(dim=1) + + def _learning_conn(self, rule, w0, nu, rng=(-1.0, 1.0)): + """MCC with a single learnable Weight; returns (connection, feature).""" + src_n, tgt_n = w0.shape + conn = self._make_mcc( + [ + tf.Weight( + name="w", value=w0.clone(), learning_rule=rule, nu=nu, range=rng + ) + ], + src_n, + tgt_n, + batch=1, + ) + return conn, conn.pipeline[0] + + def _mstdp_seq(self): + return [ + (torch.tensor([1.0, 0.0, 0.0]), torch.tensor([0.0, 0.0]), 0.0), + (torch.tensor([0.0, 0.0, 0.0]), torch.tensor([1.0, 0.0]), 1.0), + (torch.tensor([0.0, 0.0, 0.0]), torch.tensor([0.0, 0.0]), 1.0), + (torch.tensor([0.0, 0.0, 0.0]), torch.tensor([0.0, 0.0]), 1.0), + ] + + def _mstdp_reference( + self, w0, seq, nu0, dt, tc_plus=20.0, tc_minus=20.0, rng=(-1.0, 1.0) + ): + w = w0.clone().float() + src_n, tgt_n = w.shape + p_plus, p_minus = torch.zeros(src_n), torch.zeros(tgt_n) + elig = torch.zeros(src_n, tgt_n) + dp, dm = math.exp(-dt / tc_plus), math.exp(-dt / tc_minus) + for src_s, tgt_s, reward in seq: + w = w + nu0 * reward * elig + p_plus = dp * p_plus + src_s + p_minus = dm * p_minus - tgt_s + elig = torch.outer(p_plus, tgt_s) + torch.outer(src_s, p_minus) + w = torch.clamp(w, rng[0], rng[1]) + return w + + def _mstdpet_reference( + self, w0, seq, nu0, dt, tc_plus=20.0, tc_minus=20.0, tc_e=25.0, rng=(-1.0, 1.0) + ): + w = w0.clone().float() + src_n, tgt_n = w.shape + p_plus, p_minus = torch.zeros(src_n), torch.zeros(tgt_n) + elig = torch.zeros(src_n, tgt_n) + elig_tr = torch.zeros(src_n, tgt_n) + for src_s, tgt_s, reward in seq: + elig_tr = elig_tr * math.exp(-dt / tc_e) + elig / tc_e + w = w + nu0 * dt * reward * elig_tr + p_plus = p_plus * math.exp(-dt / tc_plus) + src_s + p_minus = p_minus * math.exp(-dt / tc_minus) - tgt_s + elig = torch.outer(p_plus, tgt_s) + torch.outer(src_s, p_minus) + w = torch.clamp(w, rng[0], rng[1]) + return w + + # ----------------------------------------------------------------------- # + # Individual feature outputs # + # ----------------------------------------------------------------------- # + + def test_weight_feature_output(self): + s = torch.tensor([[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]]) + w = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]) + conn = self._make_mcc([tf.Weight(name="w", value=w.clone())], 3, 2, batch=2) + assert torch.allclose(conn.compute(s), s @ w, atol=1e-6) + + def test_mask_feature_output(self): + s = torch.tensor([[1.0, 1.0, 1.0]]) + w = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]) + mask = torch.tensor([[True, False], [True, True], [False, True]]) + conn = self._make_mcc( + [ + tf.Weight(name="w", value=w.clone()), + tf.Mask(name="m", value=mask.clone()), + ], + 3, + 2, + ) + assert torch.allclose(conn.compute(s), s @ (w * mask.float()), atol=1e-6) + + def test_bias_feature_output(self): + # Bias is additive per-synapse; after the source-sum it adds bias.sum(0). + s = torch.tensor([[1.0, 0.0, 1.0], [1.0, 1.0, 0.0]]) + w = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]) + bias = torch.tensor([[0.5, -0.5], [0.1, 0.2], [0.0, 1.0]]) + conn = self._make_mcc( + [ + tf.Weight(name="w", value=w.clone()), + tf.Bias(name="b", value=bias.clone()), + ], + 3, + 2, + batch=2, + ) + assert torch.allclose(conn.compute(s), s @ w + bias.sum(0), atol=1e-6) + + def test_intensity_feature_output(self): + # Intensity's value is a per-synapse [src, tgt] tensor (a constant 2.0 here). + s = torch.tensor([[1.0, 0.0, 1.0]]) + w = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]) + intensity = torch.full((3, 2), 2.0) + conn = self._make_mcc( + [ + tf.Weight(name="w", value=w.clone()), + tf.Intensity(name="i", value=intensity.clone(), range=(-5.0, 5.0)), + ], + 3, + 2, + ) + assert torch.allclose(conn.compute(s), s @ (w * intensity), atol=1e-6) + + def test_degradation_feature_output(self): + # Degradation subtracts degrade_function(value) per-synapse -> -sum(0). + s = torch.tensor([[1.0, 1.0, 0.0]]) + w = torch.tensor([[0.4, 0.2], [0.3, 0.4], [0.5, 0.6]]) + deg = torch.tensor([[0.2, 0.4], [0.6, 0.8], [0.1, 0.3]]) + conn = self._make_mcc( + [ + tf.Weight(name="w", value=w.clone()), + tf.Degradation( + name="d", value=deg.clone(), degrade_function=lambda v: v * 0.5 + ), + ], + 3, + 2, + ) + assert torch.allclose(conn.compute(s), s @ w - (0.5 * deg).sum(0), atol=1e-6) + + def test_probability_feature_deterministic_bounds(self): + # bernoulli(1) == 1 (always passes); bernoulli(0) == 0 (always blocked). + s = torch.tensor([[1.0, 0.0, 1.0], [1.0, 1.0, 1.0]]) + w = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]) + passes = self._make_mcc( + [ + tf.Probability(name="p", value=torch.ones(3, 2)), + tf.Weight(name="w", value=w.clone()), + ], + 3, + 2, + batch=2, + ) + assert torch.allclose(passes.compute(s), s @ w, atol=1e-6) + blocked = self._make_mcc( + [ + tf.Probability(name="p", value=torch.zeros(3, 2)), + tf.Weight(name="w", value=w.clone()), + ], + 3, + 2, + batch=2, + ) + assert torch.allclose(blocked.compute(s), torch.zeros(2, 2), atol=1e-6) + + def test_adaptation_features_output(self): + for cls in (tf.AdaptationBaseSynapsHistory, tf.AdaptationBaseOtherSynaps): + src_n, tgt_n = 3, 2 + s = torch.tensor([[1.0, 0.0, 1.0]]) + feat = cls( + name="a", + value=torch.zeros(src_n, tgt_n), + ann_values=[torch.zeros(1, 1), torch.zeros(1, 1)], + ) + conn = self._make_mcc([feat], src_n, tgt_n, batch=1) + out = conn.compute(s) + assert torch.allclose(feat.value.float(), torch.ones(src_n, tgt_n)) + assert torch.allclose(out, s @ torch.ones(src_n, tgt_n), atol=1e-6) + + # ----------------------------------------------------------------------- # + # Combined pipelines == pre-collapse expansion # + # ----------------------------------------------------------------------- # + + def test_pipeline_matches_expansion(self): + torch.manual_seed(0) + src_n, tgt_n, batch = 4, 3, 2 + s = (torch.rand(batch, src_n) > 0.4).float() + w = torch.randn(src_n, tgt_n) + w2 = torch.randn(src_n, tgt_n) + mask = torch.rand(src_n, tgt_n) > 0.5 + bias = torch.randn(src_n, tgt_n) * 0.2 + deg = torch.rand(src_n, tgt_n) + + pipelines = { + "weight": [tf.Weight(name="w", value=w.clone())], + "weight+mask": [ + tf.Weight(name="w", value=w.clone()), + tf.Mask(name="m", value=mask.clone()), + ], + "weight+bias": [ + tf.Weight(name="w", value=w.clone()), + tf.Bias(name="b", value=bias.clone()), + ], + "weight+bias+degradation": [ + tf.Weight(name="w", value=w.clone()), + tf.Bias(name="b", value=bias.clone()), + tf.Degradation( + name="d", value=deg.clone(), degrade_function=lambda v: v * 0.3 + ), + ], + "weight+intensity+bias": [ + tf.Weight(name="w", value=w.clone()), + tf.Intensity( + name="i", value=torch.full((src_n, tgt_n), 1.5), range=(-5.0, 5.0) + ), + tf.Bias(name="b", value=bias.clone()), + ], + "weight,bias,weight,bias": [ + tf.Weight(name="w", value=w.clone()), + tf.Bias(name="b", value=bias.clone()), + tf.Weight(name="w2", value=w2.clone()), + tf.Bias(name="b2", value=(bias * 0.5).clone()), + ], + } + + for name, pipe in pipelines.items(): + conn = self._make_mcc(pipe, src_n, tgt_n, batch=batch) + out = conn.compute(s) + ref = self._reference_expansion(pipe, s, tgt_n) + assert torch.allclose( + out, ref, atol=1e-5 + ), f"{name}: {(out - ref).abs().max()}" + + # ----------------------------------------------------------------------- # + # Sparse activity / sparse_compute == dense # + # ----------------------------------------------------------------------- # + + def test_sparse_compute_matches_dense(self): + torch.manual_seed(1) + src_n, tgt_n, batch = 6, 4, 3 + w = torch.randn(src_n, tgt_n) + bias = torch.randn(src_n, tgt_n) * 0.2 + deg = torch.rand(src_n, tgt_n) + + def pipes(): + return [ + [tf.Weight(name="w", value=w.clone())], + [ + tf.Weight(name="w", value=w.clone()), + tf.Bias(name="b", value=bias.clone()), + ], + [ + tf.Weight(name="w", value=w.clone()), + tf.Degradation( + name="d", value=deg.clone(), degrade_function=lambda v: v * 0.4 + ), + ], + ] + + spike_sets = [ + (torch.rand(batch, src_n) > 0.5).float(), + torch.zeros(batch, src_n), # empty-spike edge case + ] + for dense_pipe, sparse_pipe in zip(pipes(), pipes()): + for s in spike_sets: + dense = self._make_mcc(dense_pipe, src_n, tgt_n, batch=batch).compute(s) + sparse = self._make_mcc( + sparse_pipe, src_n, tgt_n, batch=batch, sparse_compute=True + ).compute(s) + assert torch.allclose(dense, sparse, atol=1e-5) + + # ----------------------------------------------------------------------- # + # Regressions # + # ----------------------------------------------------------------------- # + + def test_sparse_probability_feature(self): + # Probability(sparse=True) stores its value as a sparse [1, src, tgt] + # tensor; the fold must densify it to a [src, tgt] factor (regression: + # "expand is unsupported for Sparse tensors"). + s = torch.tensor([[1.0, 0.0, 1.0], [1.0, 1.0, 1.0]]) + w = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]) + for sparse_compute in (False, True): + passes = self._make_mcc( + [ + tf.Weight(name="w", value=w.clone()), + tf.Probability(name="p", value=torch.ones(3, 2), sparse=True), + ], + 3, + 2, + batch=2, + sparse_compute=sparse_compute, + ) + assert torch.allclose(passes.compute(s), s @ w, atol=1e-6) + blocked = self._make_mcc( + [ + tf.Weight(name="w", value=w.clone()), + tf.Probability(name="p", value=torch.zeros(3, 2), sparse=True), + ], + 3, + 2, + batch=2, + sparse_compute=sparse_compute, + ) + assert torch.allclose(blocked.compute(s), torch.zeros(2, 2), atol=1e-6) + + def test_empty_pipeline_fan_in(self): + # No features: every source contributes with unit weight. + conn = self._make_mcc([], 5, 3, batch=2) + s = torch.tensor([[1.0, 1.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 0.0, 0.0]]) + out = conn.compute(s) + assert torch.allclose(out, s.sum(1, keepdim=True).expand(2, 3)) + + def test_subfeature_folds_as_identity(self): + # A sub-feature runs its side effect and contributes nothing to the + # fold, even when it precedes every real feature (regression: returned + # the int 1, which broke torch.is_floating_point when first). + w = torch.tensor([[0.5, 1.5], [1.0, 0.5], [0.5, 1.0]]) + conn = self._make_mcc([tf.Weight(name="w", value=w.clone(), norm=2.0)], 3, 2) + wf = conn.pipeline[0] + conn.pipeline = [tf.Normalization(name="n", parent_feature=wf), wf] + s = torch.tensor([[1.0, 1.0, 1.0]]) + out = conn.compute(s) + # normalize ran inside the fold (before Weight), so each target column + # of the weight sums to `norm` and the output uses those values. + assert torch.allclose(wf.value.sum(0), torch.full((2,), 2.0), atol=1e-5) + assert torch.allclose(out, s @ wf.value, atol=1e-5) + + def test_mstdp_decay_tracks_dt_change(self): + # The cached MSTDP trace-decay factors must follow connection.dt + # (regression: frozen at first update). + w0 = torch.rand(3, 2) + conn, feat = self._learning_conn(mcc.MSTDP, w0, nu=(0.1, 0.1)) + rule = feat.learning_rule + conn.source.s = torch.tensor([[1.0, 0.0, 0.0]]) + conn.target.s = torch.tensor([[0.0, 1.0]]) + rule.update(reward=0.0) + assert torch.allclose(rule._decay_plus, torch.exp(torch.tensor(-1.0 / 20.0))) + conn.dt = 5.0 + rule.update(reward=0.0) + assert torch.allclose(rule._decay_plus, torch.exp(torch.tensor(-5.0 / 20.0))) + + def test_sparse_compute_matches_dense_cuda(self): + # On CUDA the gather is gated by connection size; both the gated-off + # (small) and gated-on (large) paths must match the dense result. + if not torch.cuda.is_available(): + return + torch.manual_seed(2) + dev = torch.device("cuda") + for src_n, tgt_n in ((80, 40), (2100, 2000)): # below / above the gate + w = torch.randn(src_n, tgt_n, device=dev) + s = (torch.rand(2, src_n, device=dev) > 0.9).float() + outs = [] + for sc in (False, True): + conn = MulticompartmentConnection( + source=Input(n=src_n), + target=LIFNodes(n=tgt_n), + device=dev, + pipeline=[tf.Weight(name="w", value=w.clone())], + sparse_compute=sc, + ) + outs.append(conn.compute(s)) + assert torch.allclose(outs[0], outs[1], atol=1e-4) + + # ----------------------------------------------------------------------- # + # Performance-path equivalence # + # ----------------------------------------------------------------------- # + + def _run_mstdp(self, batch, reduction, reward_seq, seed=0): + """Run an MSTDP-learned Weight over a spike/reward sequence.""" + torch.manual_seed(seed) + src_n, tgt_n = 7, 5 + w0 = torch.rand(src_n, tgt_n) + conn = self._make_mcc( + [ + tf.Weight( + name="w", + value=w0.clone(), + learning_rule=mcc.MSTDP, + nu=(0.05, 0.05), + range=[-10, 10], + reduction=reduction, + ) + ], + src_n, + tgt_n, + batch=batch, + ) + feat = conn.pipeline[0] + rule = feat.learning_rule + torch.manual_seed(seed + 1) + for r in reward_seq: + conn.source.s = torch.bernoulli(torch.full((batch, src_n), 0.4)) + conn.target.s = torch.bernoulli(torch.full((batch, tgt_n), 0.4)) + rule.update(reward=r) + return feat.value.clone() + + def test_mstdp_rank1_matches_dense(self): + # The rank-1 addmm_ fast path (default reductions) must match the + # dense-eligibility path (forced here via equivalent custom lambdas). + rewards = [0.0, 1.0, 0.5, -2.0, 1.0, 0.0, 3.0] + slow_squeeze = lambda x, dim: torch.squeeze(x, dim) + slow_sum = lambda x, dim: torch.sum(x, dim) + for batch, fast_red, slow_red in ( + (1, None, slow_squeeze), + (4, torch.sum, slow_sum), + ): + w_fast = self._run_mstdp(batch, fast_red, rewards) + w_slow = self._run_mstdp(batch, slow_red, rewards) + assert torch.allclose( + w_fast, w_slow, atol=1e-5 + ), f"batch={batch}: {(w_fast - w_slow).abs().max()}" + # Tensor rewards take the sync-free tensor branch; same numbers. + w_fast = self._run_mstdp(1, None, [torch.tensor(r) for r in rewards]) + w_slow = self._run_mstdp(1, slow_squeeze, rewards) + assert torch.allclose(w_fast, w_slow, atol=1e-5) + + def test_fold_cache_static_and_invalidation(self): + # Static pipelines cache the folded factors; dynamic ones must not; + # learning updates through the connection invalidate the cache. + s = torch.ones(1, 6, dtype=torch.bool) + w = torch.rand(6, 4) + static = self._make_mcc([tf.Weight(name="w", value=w.clone())], 6, 4) + out1 = static.compute(s) + assert static._fold_cache is not None + assert torch.allclose(static.compute(s), out1) + + dynamic = self._make_mcc( + [ + tf.Weight(name="w", value=w.clone()), + tf.Probability(name="p", value=torch.full((6, 4), 0.5)), + ], + 6, + 4, + ) + dynamic.compute(s) + assert dynamic._fold_cache is None + + # A learning step through connection.update must drop the cache and + # the next compute must see the new weights. + learned = self._make_mcc( + [ + tf.Weight( + name="w", + value=w.clone(), + learning_rule=mcc.PostPre, + nu=(0.5, 0.5), + range=[-10, 10], + ) + ], + 6, + 4, + ) + learned.compute(s) + learned.source.s = torch.ones(1, 6) + learned.target.s = torch.ones(1, 4) + learned.source.x = torch.ones(1, 6) + learned.target.x = torch.ones(1, 4) + learned.update(learning=True) + assert learned._fold_cache is None + w_new = learned.pipeline[0].value + assert torch.allclose(learned.compute(s), s.float() @ w_new, atol=1e-5) + + def test_time_step_norm_deferred(self): + # Per-time-step normalization runs after the fold: each step's output + # uses the pre-normalization weights (old expansion semantics). + w0 = torch.tensor([[1.0, 4.0], [3.0, 4.0], [1.0, 2.0]]) + conn = self._make_mcc( + [ + tf.Weight( + name="w", value=w0.clone(), norm=1.0, norm_frequency="time step" + ) + ], + 3, + 2, + ) + s = torch.ones(1, 3) + ref_w = w0.clone() + for step in range(3): + out = conn.compute(s) + assert torch.allclose(out, s @ ref_w, atol=1e-5), f"step {step}" + ref_w = ref_w / ref_w.abs().sum(0, keepdim=True) + assert conn._fold_cache is None # never cached while norm runs per step + + # ----------------------------------------------------------------------- # + # MCC learning rules # + # ----------------------------------------------------------------------- # + + def test_noop_leaves_weight_unchanged(self): + w0 = torch.rand(3, 2) + conn, feat = self._learning_conn(mcc.NoOp, w0, nu=(0.1, 0.1)) + conn.source.s = torch.ones(1, 3) + conn.target.s = torch.ones(1, 2) + feat.learning_rule.update(reward=1.0) + assert torch.allclose(feat.value, w0) + + def test_postpre_predictable(self): + # PostPre: dW = -nu0 * outer(src_s, tgt_x) + nu1 * outer(src_x, tgt_s). + w0 = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.0, -0.2]]) + conn, feat = self._learning_conn(mcc.PostPre, w0, nu=(0.1, 0.2)) + conn.source.s = torch.tensor([[1.0, 0.0, 1.0]]) + conn.target.s = torch.tensor([[1.0, 1.0]]) + conn.source.x = torch.tensor([[0.3, 0.7, 0.2]]) + conn.target.x = torch.tensor([[0.5, 0.4]]) + feat.learning_rule.update() + dw = -0.1 * torch.outer(conn.source.s[0], conn.target.x[0]) + 0.2 * torch.outer( + conn.source.x[0], conn.target.s[0] + ) + expected = torch.clamp(w0 + dw, -1.0, 1.0) + assert torch.allclose(feat.value, expected, atol=1e-6) + assert not torch.allclose(feat.value, w0) # sanity: it actually changed + + def test_learning_respects_range_clamp(self): + # Post-only potentiation of +5 per synapse must clamp to the range max (1.0). + w0 = torch.full((2, 2), 0.9) + conn, feat = self._learning_conn( + mcc.PostPre, w0, nu=(0.0, 5.0), rng=(-1.0, 1.0) + ) + conn.source.s = torch.zeros(1, 2) + conn.target.s = torch.ones(1, 2) + conn.source.x = torch.ones(1, 2) + conn.target.x = torch.zeros(1, 2) + feat.learning_rule.update() + assert torch.allclose(feat.value, torch.ones(2, 2)) + + def test_mstdp_predictable(self): + w0 = torch.zeros(3, 2) + conn, feat = self._learning_conn(mcc.MSTDP, w0, nu=(0.5, 0.0)) + seq = self._mstdp_seq() + for src_v, tgt_v, reward in seq: + conn.source.s = src_v.unsqueeze(0) + conn.target.s = tgt_v.unsqueeze(0) + feat.learning_rule.update(reward=reward) + expected = self._mstdp_reference(w0, seq, nu0=0.5, dt=1.0) + assert torch.allclose(feat.value, expected, atol=1e-5) + assert not torch.allclose(feat.value, w0) + + def test_mstdpet_predictable(self): + w0 = torch.zeros(3, 2) + conn, feat = self._learning_conn(mcc.MSTDPET, w0, nu=(0.5, 0.5)) + seq = self._mstdp_seq() + for src_v, tgt_v, reward in seq: + conn.source.s = src_v.unsqueeze(0) + conn.target.s = tgt_v.unsqueeze(0) + feat.learning_rule.update(reward=reward) + expected = self._mstdpet_reference(w0, seq, nu0=0.5, dt=1.0) + assert torch.allclose(feat.value, expected, atol=1e-5) + assert not torch.allclose(feat.value, w0) + + def test_hebbian_predictable(self): + # Hebbian: dW = nu0 * outer(src_s, tgt_x) + nu1 * outer(src_x, tgt_s), then + # clamped. Both pre- and post-synaptic terms are positive (contrast PostPre, + # which subtracts the pre term). + w0 = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.0, -0.2]]) + conn, feat = self._learning_conn(mcc.Hebbian, w0, nu=(0.1, 0.2)) + conn.source.s = torch.tensor([[1.0, 0.0, 1.0]]) + conn.target.s = torch.tensor([[1.0, 1.0]]) + conn.source.x = torch.tensor([[0.3, 0.7, 0.2]]) + conn.target.x = torch.tensor([[0.5, 0.4]]) + feat.learning_rule.update() + dw = 0.1 * torch.outer(conn.source.s[0], conn.target.x[0]) + 0.2 * torch.outer( + conn.source.x[0], conn.target.s[0] + ) + expected = torch.clamp(w0 + dw, -1.0, 1.0) + assert torch.allclose(feat.value, expected, atol=1e-6) + assert not torch.allclose(feat.value, w0) # sanity: it actually changed + + if __name__ == "__main__": + # MulticompartmentConnection + MCC learning-rule tests (discovered + # dynamically so this list cannot go stale). + mcc_tester = TestMultiCompartmentConnection() + mcc_tests = [ + getattr(mcc_tester, n) for n in sorted(dir(mcc_tester)) if n.startswith("test_") + ] + for mcc_test in mcc_tests: + mcc_test() + print(f" PASSED: {mcc_test.__name__}") + print(f"All {len(mcc_tests)} MulticompartmentConnection tests passed.") + tester = TestConnection() # tester.test_transfer() @@ -180,7 +795,7 @@ def test_weights(self, conn_type, shape_a, shape_b, shape_w, *args, **kwargs): for update_rule in (Hebbian, PostPre, WeightDependentPostPre, MSTDP, MSTDPET, Rmax): print("Learning Rule:", update_rule) for conn_type, arg in zip(conn_types, args): - tester.test_weights(conn_type, nu=1e-2, update_rule=update_rule, *arg) + tester.check_weights(conn_type, nu=1e-2, update_rule=update_rule, *arg) # Other connections # Note: Does not include MaxPool2dConnection because this connection @@ -188,4 +803,4 @@ def test_weights(self, conn_type, shape_a, shape_b, shape_w, *args, **kwargs): conn_types = [MeanFieldConnection] args = [[[1, 28, 28], [1, 26, 26], (1, 26), 3, 1]] for conn_type, arg in zip(conn_types, args): - tester.test_weights(conn_type, decay=1, update_rule=NoOp, *arg) + tester.check_weights(conn_type, decay=1, update_rule=NoOp, *arg)