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.

[erenaydoslu]

Hi @jeshraghian,

Thank you for taking the time to review. I've pushed two new commits to address your feedback.

• Added validation in SpikeTime so that when multi_spike=False, only scalar on_target/off_target values are accepted, and when multi_spike=True, both are converted to equal-length lists. This should prevent the torch.full argument mismatch you observed.
• I refactored the Tolerance.forward method, as this method resulted in dtype mismatch errors when tensors were a mix of Long/Float dtypes. The function is also easier to read now.

I've run the new versions, and the new loss now handles both scalars and multi-spike targets without errors. Please let me know if there's anything else I can adjust.