snntorch.surrogate
By default, PyTorch’s autodifferentiation tools are unable to calculate the analytical derivative of the spiking neuron graph.
The discrete nature of spikes makes it difficult for torch.autograd
to calculate a gradient that facilitates learning.
snntorch
overrides the default gradient by using snntorch.surrogate.ATan
.
Alternative gradients are also available in the snntorch.surrogate
module.
These represent either approximations of the backward pass or probabilistic models of firing as a function of the membrane potential.
Custom, user-defined surrogate gradients can also be implemented.
At present, the surrogate gradient functions available include:
amongst several other options.
For further reading, see:
E. O. Neftci, H. Mostafa, F. Zenke (2019) Surrogate Gradient Learning in Spiking Neural Networks: Bringing the Power of Gradient-Based Optimization to Spiking Neural Networks. IEEE Signal Processing Magazine, pp. 51-63.
How to use surrogate
The surrogate gradient must be passed as the spike_grad
argument to the neuron model.
If spike_grad
is left unspecified, it defaults to snntorch.surrogate.ATan
.
In the following example, we apply the fast sigmoid surrogate to snntorch.Synaptic
.
Example:
import snntorch as snn
from snntorch import surrogate
import torch
import torch.nn as nn
alpha = 0.9
beta = 0.85
# Initialize surrogate gradient
spike_grad1 = surrogate.fast_sigmoid() # passes default parameters from a closure
spike_grad2 = surrogate.FastSigmoid.apply # passes default parameters, equivalent to above
spike_grad3 = surrogate.fast_sigmoid(slope=50) # custom parameters from a closure
# Define Network
class Net(nn.Module):
def __init__(self):
super().__init__()
# Initialize layers, specify the ``spike_grad`` argument
self.fc1 = nn.Linear(num_inputs, num_hidden)
self.lif1 = snn.Synaptic(alpha=alpha, beta=beta, spike_grad=spike_grad1)
self.fc2 = nn.Linear(num_hidden, num_outputs)
self.lif2 = snn.Synaptic(alpha=alpha, beta=beta, spike_grad=spike_grad3)
def forward(self, x, syn1, mem1, spk1, syn2, mem2):
cur1 = self.fc1(x)
spk1, syn1, mem1 = self.lif1(cur1, syn1, mem1)
cur2 = self.fc2(spk1)
spk2, syn2, mem2 = self.lif2(cur2, syn2, mem2)
return syn1, mem1, spk1, syn2, mem2, spk2
net = Net().to(device)
Custom Surrogate Gradients
For flexibility, custom surrogate gradients can also be defined by the user using custom_surrogate.
Example:
import snntorch as snn
from snntorch import surrogate
import torch
import torch.nn as nn
beta = 0.9
# Define custom surrogate gradient
def custom_fast_sigmoid(input_, grad_input, spikes):
## The hyperparameter slope is defined inside the function.
slope = 25
grad = grad_input / (slope * torch.abs(input_) + 1.0) ** 2
return grad
spike_grad = surrogate.custom_surrogate(custom_fast_sigmoid)
# Define Network
class Net(nn.Module):
def __init__(self):
super().__init__()
# Initialize layers
self.fc1 = nn.Linear(num_inputs, num_hidden)
self.lif1 = snn.Leaky(beta=beta, spike_grad=spike_grad)
self.fc2 = nn.Linear(num_hidden, num_outputs)
self.lif2 = snn.Leaky(beta=beta, spike_grad=spike_grad)
def forward(self, x, mem1, spk1, mem2):
cur1 = self.fc1(x)
spk1, mem1 = self.lif1(cur1, mem1)
cur2 = self.fc2(spk1)
spk2, mem2 = self.lif2(cur2, mem2)
return mem1, spk1, mem2, spk2
net = Net().to(device)
List of surrogate gradients
- class snntorch.surrogate.ATan(*args, **kwargs)[source]
Bases:
Function
Surrogate gradient of the Heaviside step function.
Forward pass: Heaviside step function shifted.
\[\begin{split}S=\begin{cases} 1 & \text{if U ≥ U$_{\rm thr}$} \\ 0 & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]Backward pass: Gradient of shifted arc-tan function.
\[\begin{split}S&≈\frac{1}{π}\text{arctan}(πU \frac{α}{2}) \\ \frac{∂S}{∂U}&=\frac{1}{π}\frac{1}{(1+(πU\frac{α}{2})^2)}\end{split}\]α defaults to 2, and can be modified by calling
surrogate.atan(alpha=2)
.Adapted from:
W. Fang, Z. Yu, Y. Chen, T. Masquelier, T. Huang, Y. Tian (2021) Incorporating Learnable Membrane Time Constants to Enhance Learning of Spiking Neural Networks. Proc. IEEE/CVF Int. Conf. Computer Vision (ICCV), pp. 2661-2671.
- static backward(ctx, grad_output)[source]
Define a formula for differentiating the operation with backward mode automatic differentiation.
This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the
vjp
function.)It must accept a context
ctx
as the first argument, followed by as many outputs as theforward()
returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs toforward()
. Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.The context can be used to retrieve tensors saved during the forward pass. It also has an attribute
ctx.needs_input_grad
as a tuple of booleans representing whether each input needs gradient. E.g.,backward()
will havectx.needs_input_grad[0] = True
if the first input toforward()
needs gradient computed w.r.t. the output.
- static forward(ctx, input_, alpha)[source]
Define the forward of the custom autograd Function.
This function is to be overridden by all subclasses. There are two ways to define forward:
Usage 1 (Combined forward and ctx):
@staticmethod def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any: pass
It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See combining-forward-context for more details
Usage 2 (Separate forward and ctx):
@staticmethod def forward(*args: Any, **kwargs: Any) -> Any: pass @staticmethod def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None: pass
The forward no longer accepts a ctx argument.
Instead, you must also override the
torch.autograd.Function.setup_context()
staticmethod to handle setting up thectx
object.output
is the output of the forward,inputs
are a Tuple of inputs to the forward.See extending-autograd for more details
The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with
ctx.save_for_backward()
if they are intended to be used inbackward
(equivalently,vjp
) orctx.save_for_forward()
if they are intended to be used for injvp
.
- class snntorch.surrogate.CustomSurrogate(*args, **kwargs)[source]
Bases:
Function
Surrogate gradient of the Heaviside step function.
Forward pass: Spike operator function.
\[\begin{split}S=\begin{cases} \frac{U(t)}{U} & \text{if U ≥ U$_{\rm thr}$} \\ 0 & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]Backward pass: User-defined custom surrogate gradient function.
The user defines the custom surrogate gradient in a separate function. It is passed in the forward static method and used in the backward static method.
The arguments of the custom surrogate gradient function are always the input of the forward pass (input_), the gradient of the input (grad_input) and the output of the forward pass (out).
** Important Note: The hyperparameters of the custom surrogate gradient function have to be defined inside of the function itself. **
Example:
import torch import torch.nn as nn import snntorch as snn from snntorch import surrogate def custom_fast_sigmoid(input_, grad_input, spikes): ## The hyperparameter slope is defined inside the function. slope = 25 grad = grad_input / (slope * torch.abs(input_) + 1.0) ** 2 return grad spike_grad = surrogate.custom_surrogate(custom_fast_sigmoid) net_seq = nn.Sequential(nn.Conv2d(1, 12, 5), nn.MaxPool2d(2), snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True), nn.Conv2d(12, 64, 5), nn.MaxPool2d(2), snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True), nn.Flatten(), nn.Linear(64*4*4, 10), snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True) ).to(device)
- static backward(ctx, grad_output)[source]
Define a formula for differentiating the operation with backward mode automatic differentiation.
This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the
vjp
function.)It must accept a context
ctx
as the first argument, followed by as many outputs as theforward()
returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs toforward()
. Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.The context can be used to retrieve tensors saved during the forward pass. It also has an attribute
ctx.needs_input_grad
as a tuple of booleans representing whether each input needs gradient. E.g.,backward()
will havectx.needs_input_grad[0] = True
if the first input toforward()
needs gradient computed w.r.t. the output.
- static forward(ctx, input_, custom_surrogate_function)[source]
Define the forward of the custom autograd Function.
This function is to be overridden by all subclasses. There are two ways to define forward:
Usage 1 (Combined forward and ctx):
@staticmethod def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any: pass
It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See combining-forward-context for more details
Usage 2 (Separate forward and ctx):
@staticmethod def forward(*args: Any, **kwargs: Any) -> Any: pass @staticmethod def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None: pass
The forward no longer accepts a ctx argument.
Instead, you must also override the
torch.autograd.Function.setup_context()
staticmethod to handle setting up thectx
object.output
is the output of the forward,inputs
are a Tuple of inputs to the forward.See extending-autograd for more details
The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with
ctx.save_for_backward()
if they are intended to be used inbackward
(equivalently,vjp
) orctx.save_for_forward()
if they are intended to be used for injvp
.
- class snntorch.surrogate.FastSigmoid(*args, **kwargs)[source]
Bases:
Function
Surrogate gradient of the Heaviside step function.
Forward pass: Heaviside step function shifted.
\[\begin{split}S=\begin{cases} 1 & \text{if U ≥ U$_{\rm thr}$} \\ 0 & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]Backward pass: Gradient of fast sigmoid function.
\[\begin{split}S&≈\frac{U}{1 + k|U|} \\ \frac{∂S}{∂U}&=\frac{1}{(1+k|U|)^2}\end{split}\]\(k\) defaults to 25, and can be modified by calling
surrogate.fast_sigmoid(slope=25)
.Adapted from:
F. Zenke, S. Ganguli (2018) SuperSpike: Supervised Learning in Multilayer Spiking Neural Networks. Neural Computation, pp. 1514-1541.
- static backward(ctx, grad_output)[source]
Define a formula for differentiating the operation with backward mode automatic differentiation.
This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the
vjp
function.)It must accept a context
ctx
as the first argument, followed by as many outputs as theforward()
returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs toforward()
. Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.The context can be used to retrieve tensors saved during the forward pass. It also has an attribute
ctx.needs_input_grad
as a tuple of booleans representing whether each input needs gradient. E.g.,backward()
will havectx.needs_input_grad[0] = True
if the first input toforward()
needs gradient computed w.r.t. the output.
- static forward(ctx, input_, slope)[source]
Define the forward of the custom autograd Function.
This function is to be overridden by all subclasses. There are two ways to define forward:
Usage 1 (Combined forward and ctx):
@staticmethod def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any: pass
It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See combining-forward-context for more details
Usage 2 (Separate forward and ctx):
@staticmethod def forward(*args: Any, **kwargs: Any) -> Any: pass @staticmethod def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None: pass
The forward no longer accepts a ctx argument.
Instead, you must also override the
torch.autograd.Function.setup_context()
staticmethod to handle setting up thectx
object.output
is the output of the forward,inputs
are a Tuple of inputs to the forward.See extending-autograd for more details
The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with
ctx.save_for_backward()
if they are intended to be used inbackward
(equivalently,vjp
) orctx.save_for_forward()
if they are intended to be used for injvp
.
- snntorch.surrogate.Heaviside(*args, **kwargs)[source]
Default spiking function for neuron.
Forward pass: Heaviside step function shifted.
\[\begin{split}S=\begin{cases} 1 & \text{if U ≥ U$_{\rm thr}$} \\ 0 & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]Backward pass: Heaviside step function shifted.
\[\begin{split}\frac{∂S}{∂U}=\begin{cases} 1 & \text{if U ≥ U$_{\rm thr}$} \\ 0 & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]Although the backward pass is clearly not the analytical solution of the forward pass, this assumption holds true on the basis that a reset necessarily occurs after a spike is generated when \(U ≥ U_{\rm thr}\).
- snntorch.surrogate.LSO(slope=0.1)[source]
Leaky spike operator gradient enclosed with a parameterized slope.
- class snntorch.surrogate.LeakySpikeOperator(*args, **kwargs)[source]
Bases:
Function
Surrogate gradient of the Heaviside step function.
Forward pass: Spike operator function.
\[\begin{split}S=\begin{cases} \frac{U(t)}{U} & \text{if U ≥ U$_{\rm thr}$} \\ 0 & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]Backward pass: Leaky gradient of spike operator, where the subthreshold gradient is treated as a small constant slope.
\[\begin{split}S&≈\begin{cases} \frac{U(t)}{U}\Big{|}_{U(t)→U_{\rm thr}} & \text{if U ≥ U$_{\rm thr}$} \\ kU & \text{if U < U$_{\rm thr}$}\end{cases} \\ \frac{∂S}{∂U}&=\begin{cases} 1 & \text{if U ≥ U$_{\rm thr}$} \\ k & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]\(k\) defaults to 0.1, and can be modified by calling
surrogate.LSO(slope=0.1)
.The gradient is identical to that of a threshold-shifted Leaky ReLU.
- static backward(ctx, grad_output)[source]
Define a formula for differentiating the operation with backward mode automatic differentiation.
This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the
vjp
function.)It must accept a context
ctx
as the first argument, followed by as many outputs as theforward()
returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs toforward()
. Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.The context can be used to retrieve tensors saved during the forward pass. It also has an attribute
ctx.needs_input_grad
as a tuple of booleans representing whether each input needs gradient. E.g.,backward()
will havectx.needs_input_grad[0] = True
if the first input toforward()
needs gradient computed w.r.t. the output.
- static forward(ctx, input_, slope)[source]
Define the forward of the custom autograd Function.
This function is to be overridden by all subclasses. There are two ways to define forward:
Usage 1 (Combined forward and ctx):
@staticmethod def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any: pass
It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See combining-forward-context for more details
Usage 2 (Separate forward and ctx):
@staticmethod def forward(*args: Any, **kwargs: Any) -> Any: pass @staticmethod def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None: pass
The forward no longer accepts a ctx argument.
Instead, you must also override the
torch.autograd.Function.setup_context()
staticmethod to handle setting up thectx
object.output
is the output of the forward,inputs
are a Tuple of inputs to the forward.See extending-autograd for more details
The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with
ctx.save_for_backward()
if they are intended to be used inbackward
(equivalently,vjp
) orctx.save_for_forward()
if they are intended to be used for injvp
.
- snntorch.surrogate.SFS(slope=25, B=1)[source]
SparseFastSigmoid surrogate gradient enclosed with a parameterized slope and sparsity threshold.
- snntorch.surrogate.SSO(mean=0, variance=0.2)[source]
Stochastic spike operator gradient enclosed with a parameterized mean and variance.
- class snntorch.surrogate.Sigmoid(*args, **kwargs)[source]
Bases:
Function
Surrogate gradient of the Heaviside step function.
Forward pass: Heaviside step function shifted.
\[\begin{split}S=\begin{cases} 1 & \text{if U ≥ U$_{\rm thr}$} \\ 0 & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]Backward pass: Gradient of sigmoid function.
\[\begin{split}S&≈\frac{1}{1 + {\rm exp}(-kU)} \\ \frac{∂S}{∂U}&=\frac{k {\rm exp}(-kU)}{[{\rm exp}(-kU)+1]^2}\end{split}\]\(k\) defaults to 25, and can be modified by calling
surrogate.sigmoid(slope=25)
.Adapted from:
F. Zenke, S. Ganguli (2018) SuperSpike: Supervised Learning in Multilayer Spiking Neural Networks. Neural Computation, pp. 1514-1541.
- static backward(ctx, grad_output)[source]
Define a formula for differentiating the operation with backward mode automatic differentiation.
This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the
vjp
function.)It must accept a context
ctx
as the first argument, followed by as many outputs as theforward()
returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs toforward()
. Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.The context can be used to retrieve tensors saved during the forward pass. It also has an attribute
ctx.needs_input_grad
as a tuple of booleans representing whether each input needs gradient. E.g.,backward()
will havectx.needs_input_grad[0] = True
if the first input toforward()
needs gradient computed w.r.t. the output.
- static forward(ctx, input_, slope)[source]
Define the forward of the custom autograd Function.
This function is to be overridden by all subclasses. There are two ways to define forward:
Usage 1 (Combined forward and ctx):
@staticmethod def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any: pass
It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See combining-forward-context for more details
Usage 2 (Separate forward and ctx):
@staticmethod def forward(*args: Any, **kwargs: Any) -> Any: pass @staticmethod def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None: pass
The forward no longer accepts a ctx argument.
Instead, you must also override the
torch.autograd.Function.setup_context()
staticmethod to handle setting up thectx
object.output
is the output of the forward,inputs
are a Tuple of inputs to the forward.See extending-autograd for more details
The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with
ctx.save_for_backward()
if they are intended to be used inbackward
(equivalently,vjp
) orctx.save_for_forward()
if they are intended to be used for injvp
.
- class snntorch.surrogate.SparseFastSigmoid(*args, **kwargs)[source]
Bases:
Function
Surrogate gradient of the Heaviside step function.
Forward pass: Heaviside step function shifted.
\[\begin{split}S=\begin{cases} 1 & \text{if U ≥ U$_{\rm thr}$} \\ 0 & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]Backward pass: Gradient of fast sigmoid function clipped below B.
\[\begin{split}S&≈\frac{U}{1 + k|U|}H(U-B) \\ \frac{∂S}{∂U}&=\begin{cases} \frac{1}{(1+k|U|)^2} & \text{\rm if U > B} 0 & \text{\rm otherwise} \end{cases}\end{split}\]\(k\) defaults to 25, and can be modified by calling
surrogate.SFS(slope=25)
. \(B\) defaults to 1, and can be modified by callingsurrogate.SFS(B=1)
.Adapted from:
N. Perez-Nieves and D.F.M. Goodman (2021) Sparse Spiking Gradient Descent. https://arxiv.org/pdf/2105.08810.pdf.
- static backward(ctx, grad_output)[source]
Define a formula for differentiating the operation with backward mode automatic differentiation.
This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the
vjp
function.)It must accept a context
ctx
as the first argument, followed by as many outputs as theforward()
returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs toforward()
. Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.The context can be used to retrieve tensors saved during the forward pass. It also has an attribute
ctx.needs_input_grad
as a tuple of booleans representing whether each input needs gradient. E.g.,backward()
will havectx.needs_input_grad[0] = True
if the first input toforward()
needs gradient computed w.r.t. the output.
- static forward(ctx, input_, slope, B)[source]
Define the forward of the custom autograd Function.
This function is to be overridden by all subclasses. There are two ways to define forward:
Usage 1 (Combined forward and ctx):
@staticmethod def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any: pass
It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See combining-forward-context for more details
Usage 2 (Separate forward and ctx):
@staticmethod def forward(*args: Any, **kwargs: Any) -> Any: pass @staticmethod def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None: pass
The forward no longer accepts a ctx argument.
Instead, you must also override the
torch.autograd.Function.setup_context()
staticmethod to handle setting up thectx
object.output
is the output of the forward,inputs
are a Tuple of inputs to the forward.See extending-autograd for more details
The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with
ctx.save_for_backward()
if they are intended to be used inbackward
(equivalently,vjp
) orctx.save_for_forward()
if they are intended to be used for injvp
.
- class snntorch.surrogate.SpikeRateEscape(*args, **kwargs)[source]
Bases:
Function
Surrogate gradient of the Heaviside step function.
Forward pass: Heaviside step function shifted.
\[\begin{split}S=\begin{cases} 1 & \text{if U ≥ U$_{\rm thr}$} \\ 0 & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]Backward pass: Gradient of Boltzmann Distribution.
\[\frac{∂S}{∂U}=k{\rm exp}(-β|U-1|)\]\(β\) defaults to 1, and can be modified by calling
surrogate.spike_rate_escape(beta=1)
. \(k\) defaults to 25, and can be modified by callingsurrogate.spike_rate_escape(slope=25)
.Adapted from:
Wulfram Gerstner and Werner M. Kistler,
Spiking neuron models: Single neurons, populations, plasticity. Cambridge University Press, 2002.*
- backward(grad_output)[source]
Define a formula for differentiating the operation with backward mode automatic differentiation.
This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the
vjp
function.)It must accept a context
ctx
as the first argument, followed by as many outputs as theforward()
returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs toforward()
. Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.The context can be used to retrieve tensors saved during the forward pass. It also has an attribute
ctx.needs_input_grad
as a tuple of booleans representing whether each input needs gradient. E.g.,backward()
will havectx.needs_input_grad[0] = True
if the first input toforward()
needs gradient computed w.r.t. the output.
- static forward(ctx, input_, beta, slope)[source]
Define the forward of the custom autograd Function.
This function is to be overridden by all subclasses. There are two ways to define forward:
Usage 1 (Combined forward and ctx):
@staticmethod def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any: pass
It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See combining-forward-context for more details
Usage 2 (Separate forward and ctx):
@staticmethod def forward(*args: Any, **kwargs: Any) -> Any: pass @staticmethod def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None: pass
The forward no longer accepts a ctx argument.
Instead, you must also override the
torch.autograd.Function.setup_context()
staticmethod to handle setting up thectx
object.output
is the output of the forward,inputs
are a Tuple of inputs to the forward.See extending-autograd for more details
The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with
ctx.save_for_backward()
if they are intended to be used inbackward
(equivalently,vjp
) orctx.save_for_forward()
if they are intended to be used for injvp
.
- class snntorch.surrogate.StochasticSpikeOperator(*args, **kwargs)[source]
Bases:
Function
Surrogate gradient of the Heaviside step function.
Forward pass: Spike operator function.
\[\begin{split}S=\begin{cases} \frac{U(t)}{U} & \text{if U ≥ U$_{\rm thr}$} \\ 0 & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]Backward pass: Gradient of spike operator, where the subthreshold gradient is sampled from uniformly distributed noise on the interval \((𝒰\sim[-0.5, 0.5)+μ) σ^2\), where \(μ\) is the mean and \(σ^2\) is the variance.
\[\begin{split}S&≈\begin{cases} \frac{U(t)}{U}\Big{|}_{U(t)→U_{\rm thr}} & \text{if U ≥ U$_{\rm thr}$} \\ (𝒰\sim[-0.5, 0.5) + μ) σ^2 & \text{if U < U$_{\rm thr}$} \end{cases} \\ \frac{∂S}{∂U}&=\begin{cases} 1 & \text{if U ≥ U$_{\rm thr}$} \\ (𝒰\sim[-0.5, 0.5) + μ) σ^2 & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]\(μ\) defaults to 0, and can be modified by calling
surrogate.SSO(mean=0)
.\(σ^2\) defaults to 0.2, and can be modified by calling
surrogate.SSO(variance=0.5)
.The above defaults set the gradient to the following expression:
\[\begin{split}\frac{∂S}{∂U}&=\begin{cases} 1 & \text{if U ≥ U$_{\rm thr}$} \\ (𝒰\sim[-0.1, 0.1) & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]- static backward(ctx, grad_output)[source]
Define a formula for differentiating the operation with backward mode automatic differentiation.
This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the
vjp
function.)It must accept a context
ctx
as the first argument, followed by as many outputs as theforward()
returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs toforward()
. Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.The context can be used to retrieve tensors saved during the forward pass. It also has an attribute
ctx.needs_input_grad
as a tuple of booleans representing whether each input needs gradient. E.g.,backward()
will havectx.needs_input_grad[0] = True
if the first input toforward()
needs gradient computed w.r.t. the output.
- static forward(ctx, input_, mean, variance)[source]
Define the forward of the custom autograd Function.
This function is to be overridden by all subclasses. There are two ways to define forward:
Usage 1 (Combined forward and ctx):
@staticmethod def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any: pass
It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See combining-forward-context for more details
Usage 2 (Separate forward and ctx):
@staticmethod def forward(*args: Any, **kwargs: Any) -> Any: pass @staticmethod def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None: pass
The forward no longer accepts a ctx argument.
Instead, you must also override the
torch.autograd.Function.setup_context()
staticmethod to handle setting up thectx
object.output
is the output of the forward,inputs
are a Tuple of inputs to the forward.See extending-autograd for more details
The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with
ctx.save_for_backward()
if they are intended to be used inbackward
(equivalently,vjp
) orctx.save_for_forward()
if they are intended to be used for injvp
.
- class snntorch.surrogate.StraightThroughEstimator(*args, **kwargs)[source]
Bases:
Function
Straight Through Estimator.
Forward pass: Heaviside step function shifted.
\[\begin{split}S=\begin{cases} 1 & \text{if U ≥ U$_{\rm thr}$} \\ 0 & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]Backward pass: Gradient of fast sigmoid function.
\[\frac{∂S}{∂U}=1\]- static backward(ctx, grad_output)[source]
Define a formula for differentiating the operation with backward mode automatic differentiation.
This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the
vjp
function.)It must accept a context
ctx
as the first argument, followed by as many outputs as theforward()
returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs toforward()
. Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.The context can be used to retrieve tensors saved during the forward pass. It also has an attribute
ctx.needs_input_grad
as a tuple of booleans representing whether each input needs gradient. E.g.,backward()
will havectx.needs_input_grad[0] = True
if the first input toforward()
needs gradient computed w.r.t. the output.
- static forward(ctx, input_)[source]
Define the forward of the custom autograd Function.
This function is to be overridden by all subclasses. There are two ways to define forward:
Usage 1 (Combined forward and ctx):
@staticmethod def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any: pass
It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See combining-forward-context for more details
Usage 2 (Separate forward and ctx):
@staticmethod def forward(*args: Any, **kwargs: Any) -> Any: pass @staticmethod def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None: pass
The forward no longer accepts a ctx argument.
Instead, you must also override the
torch.autograd.Function.setup_context()
staticmethod to handle setting up thectx
object.output
is the output of the forward,inputs
are a Tuple of inputs to the forward.See extending-autograd for more details
The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with
ctx.save_for_backward()
if they are intended to be used inbackward
(equivalently,vjp
) orctx.save_for_forward()
if they are intended to be used for injvp
.
- class snntorch.surrogate.Triangular(*args, **kwargs)[source]
Bases:
Function
Triangular Surrogate Gradient.
Forward pass: Heaviside step function shifted.
\[\begin{split}S=\begin{cases} 1 & \text{if U ≥ U$_{\rm thr}$} \\ 0 & \text{if U < U$_{\rm thr}$} \end{cases}\end{split}\]Backward pass: Gradient of the triangular function.
\[\begin{split}\frac{∂S}{∂U}=\begin{cases} U_{\rm thr} & \text{if U < U$_{\rm thr}$} \\ -U_{\rm thr} & \text{if U ≥ U$_{\rm thr}$} \end{cases}\end{split}\]- static backward(ctx, grad_output)[source]
Define a formula for differentiating the operation with backward mode automatic differentiation.
This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the
vjp
function.)It must accept a context
ctx
as the first argument, followed by as many outputs as theforward()
returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs toforward()
. Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.The context can be used to retrieve tensors saved during the forward pass. It also has an attribute
ctx.needs_input_grad
as a tuple of booleans representing whether each input needs gradient. E.g.,backward()
will havectx.needs_input_grad[0] = True
if the first input toforward()
needs gradient computed w.r.t. the output.
- static forward(ctx, input_, threshold)[source]
Define the forward of the custom autograd Function.
This function is to be overridden by all subclasses. There are two ways to define forward:
Usage 1 (Combined forward and ctx):
@staticmethod def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any: pass
It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See combining-forward-context for more details
Usage 2 (Separate forward and ctx):
@staticmethod def forward(*args: Any, **kwargs: Any) -> Any: pass @staticmethod def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None: pass
The forward no longer accepts a ctx argument.
Instead, you must also override the
torch.autograd.Function.setup_context()
staticmethod to handle setting up thectx
object.output
is the output of the forward,inputs
are a Tuple of inputs to the forward.See extending-autograd for more details
The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with
ctx.save_for_backward()
if they are intended to be used inbackward
(equivalently,vjp
) orctx.save_for_forward()
if they are intended to be used for injvp
.
- snntorch.surrogate.atan(alpha=2.0)[source]
ArcTan surrogate gradient enclosed with a parameterized slope.
- snntorch.surrogate.custom_surrogate(custom_surrogate_function)[source]
Custom surrogate gradient enclosed within a wrapper.
- snntorch.surrogate.fast_sigmoid(slope=25)[source]
FastSigmoid surrogate gradient enclosed with a parameterized slope.
- snntorch.surrogate.sigmoid(slope=25)[source]
Sigmoid surrogate gradient enclosed with a parameterized slope.
- snntorch.surrogate.spike_rate_escape(beta=1, slope=25)[source]
SpikeRateEscape surrogate gradient enclosed with a parameterized slope.