Skip to content

faster LIF forward #384

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

Merged
merged 11 commits into from
Jun 29, 2020
Merged

faster LIF forward #384

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
bc43241
faster LIF forward
SimonInParis Jun 17, 2020
a77b247
default reduction is sum. Faster code when no batch is used
SimonInParis Jun 17, 2020
05f2ebe
fixed indent bug
SimonInParis Jun 17, 2020
f24f94c
changed NoBatch function location
SimonInParis Jun 21, 2020
3e44698
moved NoBatch inside class
SimonInParis Jun 22, 2020
25f7d94
squeeze!
SimonInParis Jun 22, 2020
925f780
optim for D&C nodes
SimonInParis Jun 22, 2020
b41bd64
optim for SRM0 nodes
SimonInParis Jun 22, 2020
2efc376
optim for currentLIF & IF nodes
SimonInParis Jun 22, 2020
d80d5dc
default reduction corrected
SimonInParis Jun 23, 2020
3af105b
last changes
SimonInParis Jun 26, 2020
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
10 changes: 6 additions & 4 deletions bindsnet/learning/learning.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@
)
from ..utils import im2col_indices


class LearningRule(ABC):
# language=rst
"""
Expand Down Expand Up @@ -56,9 +55,12 @@ def __init__(

# Parameter update reduction across minibatch dimension.
if reduction is None:
reduction = torch.mean

self.reduction = reduction
if self.source.batch_size == 1:
self.reduction = torch.squeeze
else:
self.reduction = torch.sum
else:
self.reduction = reduction

# Weight decay.
self.weight_decay = weight_decay
Expand Down
42 changes: 16 additions & 26 deletions bindsnet/network/nodes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -375,12 +375,10 @@ def forward(self, x: torch.Tensor) -> None:
:param x: Inputs to the layer.
"""
# Integrate input voltages.
self.v += (self.refrac_count == 0).float() * x
self.v += (self.refrac_count <= 0).float() * x

# Decrement refractory counters.
self.refrac_count = (self.refrac_count > 0).float() * (
self.refrac_count - self.dt
)
self.refrac_count -= self.dt

# Check for spiking neurons.
self.s = self.v >= self.thresh
Expand Down Expand Up @@ -509,16 +507,16 @@ def forward(self, x: torch.Tensor) -> None:
self.v = self.decay * (self.v - self.rest) + self.rest

# Integrate inputs.
self.v += (self.refrac_count == 0).float() * x

x.masked_fill_(self.refrac_count > 0, 0.0) # OPTIM 2
# Decrement refractory counters.
self.refrac_count = (self.refrac_count > 0).float() * (
self.refrac_count - self.dt
)
self.refrac_count -= self.dt # OPTIM 1

self.v += x # interlaced

# Check for spiking neurons.
self.s = self.v >= self.thresh


# Refractoriness and voltage reset.
self.refrac_count.masked_fill_(self.s, self.refrac)
self.v.masked_fill_(self.s, self.reset)
Expand Down Expand Up @@ -653,13 +651,11 @@ def forward(self, x: torch.Tensor) -> None:
self.i *= self.i_decay

# Decrement refractory counters.
self.refrac_count = (self.refrac_count > 0).float() * (
self.refrac_count - self.dt
)
self.refrac_count -= self.dt

# Integrate inputs.
self.i += x
self.v += (self.refrac_count == 0).float() * self.i
self.v += (self.refrac_count <= 0).float() * self.i

# Check for spiking neurons.
self.s = self.v >= self.thresh
Expand Down Expand Up @@ -776,7 +772,7 @@ def __init__(
"tc_decay", torch.tensor(tc_decay)
) # Time constant of neuron voltage decay.
self.register_buffer(
"decay", torch.empty_like(self.tc_decay)
"decay", torch.empty_like(self.tc_decay, dtype=torch.float32)
) # Set in compute_decays.
self.register_buffer(
"theta_plus", torch.tensor(theta_plus)
Expand Down Expand Up @@ -808,12 +804,10 @@ def forward(self, x: torch.Tensor) -> None:
self.theta *= self.theta_decay

# Integrate inputs.
self.v += (self.refrac_count == 0).float() * x
self.v += (self.refrac_count <= 0).float() * x

# Decrement refractory counters.
self.refrac_count = (self.refrac_count > 0).float() * (
self.refrac_count - self.dt
)
self.refrac_count -= self.dt

# Check for spiking neurons.
self.s = self.v >= self.thresh + self.theta
Expand Down Expand Up @@ -965,12 +959,10 @@ def forward(self, x: torch.Tensor) -> None:
self.theta *= self.theta_decay

# Integrate inputs.
self.v += (self.refrac_count == 0).float() * x
self.v += (self.refrac_count <= 0).float() * x

# Decrement refractory counters.
self.refrac_count = (self.refrac_count > 0).float() * (
self.refrac_count - self.dt
)
self.refrac_count -= self.dt

# Check for spiking neurons.
self.s = self.v >= self.thresh + self.theta
Expand Down Expand Up @@ -1298,17 +1290,15 @@ def forward(self, x: torch.Tensor) -> None:
self.v = self.decay * (self.v - self.rest) + self.rest

# Integrate inputs.
self.v += (self.refrac_count == 0).float() * self.eps_0 * x
self.v += (self.refrac_count <= 0).float() * self.eps_0 * x

# Compute (instantaneous) probabilities of spiking, clamp between 0 and 1 using exponentials.
# Also known as 'escape noise', this simulates nearby neurons.
self.rho = self.rho_0 * torch.exp((self.v - self.thresh) / self.d_thresh)
self.s_prob = 1.0 - torch.exp(-self.rho * self.dt)

# Decrement refractory counters.
self.refrac_count = (self.refrac_count > 0).float() * (
self.refrac_count - self.dt
)
self.refrac_count -= self.dt

# Check for spiking neurons (spike when probability > some random number).
self.s = torch.rand_like(self.s_prob) < self.s_prob
Expand Down