Runs a single simulation step. :param x: Inputs to the layer.
(self, x: torch.Tensor)
| 1067 | self.one_spike = one_spike # One spike per timestep. |
| 1068 | |
| 1069 | def forward(self, x: torch.Tensor) -> None: |
| 1070 | # language=rst |
| 1071 | """ |
| 1072 | Runs a single simulation step. |
| 1073 | |
| 1074 | :param x: Inputs to the layer. |
| 1075 | """ |
| 1076 | # Decay voltages and adaptive thresholds. |
| 1077 | self.v = self.decay * (self.v - self.rest) + self.rest |
| 1078 | if self.learning: |
| 1079 | self.theta *= self.theta_decay |
| 1080 | |
| 1081 | # Integrate inputs. |
| 1082 | self.v += (self.refrac_count <= 0).float() * x |
| 1083 | |
| 1084 | # Decrement refractory counters. |
| 1085 | self.refrac_count -= self.dt |
| 1086 | |
| 1087 | # Check for spiking neurons. |
| 1088 | self.s = self.v >= self.thresh + self.theta |
| 1089 | |
| 1090 | # Refractoriness, voltage reset, and adaptive thresholds. |
| 1091 | self.refrac_count.masked_fill_(self.s, self.refrac) |
| 1092 | self.v.masked_fill_(self.s, self.reset) |
| 1093 | if self.learning: |
| 1094 | self.theta += self.theta_plus * self.s.float().sum(0) |
| 1095 | |
| 1096 | # Choose only a single neuron to spike. |
| 1097 | if self.one_spike: |
| 1098 | if self.s.any(): |
| 1099 | _any = self.s.view(self.batch_size, -1).any(1) |
| 1100 | ind = torch.multinomial( |
| 1101 | self.s.float().view(self.batch_size, -1)[_any], 1 |
| 1102 | ) |
| 1103 | _any = _any.nonzero() |
| 1104 | self.s.zero_() |
| 1105 | self.s.view(self.batch_size, -1)[_any, ind] = 1 |
| 1106 | |
| 1107 | # Voltage clipping to lower bound. |
| 1108 | if self.lbound is not None: |
| 1109 | self.v.masked_fill_(self.v < self.lbound, self.lbound) |
| 1110 | |
| 1111 | super().forward(x) |
| 1112 | |
| 1113 | def reset_state_variables(self) -> None: |
| 1114 | # language=rst |