Skip to content

Performance Improvements #775

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
C-Earl wants to merge 8 commits into
base: master
Choose a base branch
from
Open

Performance Improvements #775

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
aad3361
Reduce CPU->GPU transfer every time step with learning.
C-Earl Jul 24, 2026
1b6e4a2
Foldable pipelines
C-Earl Jul 24, 2026
2428476
More files for foldable pipelines
C-Earl Jul 24, 2026
677d196
Sparse compute
C-Earl Jul 24, 2026
884010b
Folding pipeline computation
C-Earl Jul 26, 2026
48c98c0
Fixed indentation error with learning rule class
C-Earl Jul 26, 2026
6f3c5a6
Test cases and minor bug fixes
C-Earl Jul 26, 2026
a81c05d
Benchmark tests
C-Earl Jul 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
168 changes: 93 additions & 75 deletions bindsnet/learning/MCC_learning.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -304,89 +304,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):
Expand Down Expand Up @@ -503,13 +506,28 @@ def _connection_update(self, **kwargs) -> None:
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
if not hasattr(self, "_a_plus_default"):
dev = self.feature_value.device
self._a_plus_default = torch.tensor(1.0, device=dev)
self._a_minus_default = torch.tensor(-1.0, device=dev)
dt = float(self.connection.dt)
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 = torch.tensor(
kwargs.get("a_minus", -1.0), 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)
)

# Compute weight update based on the eligibility value of the past timestep.
Expand All @@ -535,9 +553,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)
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.
Expand Down
93 changes: 61 additions & 32 deletions bindsnet/network/topology.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -413,6 +413,7 @@ def __init__(
pipeline: list = [],
manual_update: bool = False,
traces: bool = False,
sparse_compute: bool = False,
**kwargs,
) -> None:
# language=rst
Expand All @@ -426,6 +427,8 @@ 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 only when few are active).
"""

super().__init__(source, target, device, pipeline, **kwargs)
Expand All @@ -434,49 +437,75 @@ def __init__(
if self.traces:
self.activity = None

self.sparse_compute = sparse_compute

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)
# 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)
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(s.size(1), b_eff.size(-1), device=s.device)
if not torch.is_floating_point(a_eff):
a_eff = a_eff.float()

if self.sparse_compute:
# 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
# Additive terms apply to every synapse regardless of spikes, so sum over
# all source rows (the closed form of the [batch, source.n, target.n]
# source-sum).
if b_eff is not None:
out = out + b_eff.sum(dim=0)

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 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
Expand Down
Loading
Loading