MSE Temporal Loss Bugfix + Optimization

[erenaydoslu]

Training a model with mse_temporal_loss would not work at all. Turns out the function spikegen.targets_convert used for converting target indices to spike times generates a tensor full of zeros, which would end up to be the target in MSE calculations. As a result, the model learns to spike all the outputs all the time.

Furthermore, training with mse_temporal_loss is also very slow. I changed the Python loops under FirstSpike to PyTorch functions without using loops. On my device, this leads to ~30x improvement in speed.

[jeshraghian]

Thanks for working through this. Training worked when specifying on_target and off_target as a single value, though it threw an error when I'd set multiple possible spike times. E.g.,

on_target = torch.tensor((5, 10))
off_target = torch.tensor((15, 20))
loss_fn = SF.mse_temporal_loss(on_target=on_target, off_target=off_target, tolerance=1)

In your update, it raised the following error:

TypeError: full() received an invalid combination of arguments - got (tuple, Tensor, device=torch.device, dtype=torch.dtype),...

whereas in the current release of snnTorch, I'd receive this error:

RuntimeError: Boolean value of Tensor with more than one value is ambiguous

I'll attempt to debug this and will add a few tests.