(gamma)
| 32 | |
| 33 | |
| 34 | def fire_function(gamma): |
| 35 | class ZIF(torch.autograd.Function): |
| 36 | @staticmethod |
| 37 | def forward(ctx, input): |
| 38 | out = (input >= 0).float() |
| 39 | ctx.save_for_backward(input) |
| 40 | return out |
| 41 | |
| 42 | @staticmethod |
| 43 | def backward(ctx, grad_output): |
| 44 | (input, ) = ctx.saved_tensors |
| 45 | grad_input = grad_output.clone() |
| 46 | tmp = (input.abs() < gamma/2).float() / gamma |
| 47 | grad_input = grad_input * tmp |
| 48 | return grad_input, None |
| 49 | |
| 50 | return ZIF.apply |
| 51 | |
| 52 | |
| 53 | class LIFSpike(nn.Module): |