brainpy.math.surrogate.slayer_grad
brainpy.math.surrogate.slayer_grad#
- brainpy.math.surrogate.slayer_grad = <brainpy._src.math.surrogate._utils.VJPCustom object>#
Spike function with the slayer surrogate gradient function.
Forward function:
\[\begin{split}g(x) = \begin{cases} 1, & x \geq 0 \\ 0, & x < 0 \\ \end{cases}\end{split}\]Backward function:
\[g'(x) = \exp(-\alpha |x|)\]>>> import brainpy as bp >>> import brainpy.math as bm >>> import matplotlib.pyplot as plt >>> bp.visualize.get_figure(1, 1, 4, 6) >>> xs = bm.linspace(-3, 3, 1000) >>> for alpha in [0.5, 1., 2., 4.]: >>> grads = bm.vector_grad(bm.surrogate.slayer_grad)(xs, alpha) >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) >>> plt.legend() >>> plt.show()
(Source code, png, hires.png, pdf)
- Parameters
x (jax.Array, Array) – The input data.
alpha (float) – Parameter to control smoothness of gradient
- Returns
out – The spiking state.
- Return type
jax.Array
References
- 1
Shrestha, S. B. & Orchard, G. Slayer: spike layer error reassignment in time. In Advances in Neural Information Processing Systems Vol. 31, 1412–1421 (NeurIPS, 2018).