STCA-SNN: Spatio-Temporal Coordinate Attention for Spiking Neural Networks

Model Architecture

• This is the official repository for paper STCA-SNN: Spatio-Temporal Coordinate Attention for Spiking Neural Networks.

Comparison with the SOTA methods on neuromorphic datasets

Prerequisites

• Python 3.9
• Pytorch 2.5.1
• Spikingjelly 0.0.0.0.14

How to use STCA in the input stage of an SNN

importtorchimporttorch.nnasnnfromspikingjelly.activation_basedimportneuron, layer, functionalfrommodule.STCAimportSpatio_Temporal_Coordinate_AttentionclassSNNInputModule(nn.Module):
def__init__(self, T=10, in_channels=2, out_channels=64):
super().__init__()
self.T=T# Using SpikingJelly's built-in layersself.conv=layer.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
self.bn=layer.BatchNorm2d(out_channels)
self.lif=neuron.LIFNode(tau=2.0, surrogate_function=neuron.surrogate.Sigmoid())
# Configure SpikingJelly modules to process temporal sequences nativelyfunctional.set_step_mode(self, step_mode='m')
# Initialize STCA# Fine-tuning the hyperparameters ratio_1 and ratio_2 may yield better accuracyself.stca=Spatio_Temporal_Coordinate_Attention(
T=self.T,
out_channels=out_channels,
ratio_1=8,
ratio_2=16
)
defforward(self, x):
# Input shape: [Batch, Time, Channels, Height, Width]# 1. Convert to SpikingJelly's expected format: [Time, Batch, C, H, W]x_t=x.transpose(0, 1)
# 2. Extract continuous potentials and discrete spikesmem_potentials_t=self.bn(self.conv(x_t))
spikes_t=self.lif(mem_potentials_t)
# 3. Revert back to STCA format: [Batch, Time, C, H, W]mem_potentials=mem_potentials_t.transpose(0, 1)
spikes=spikes_t.transpose(0, 1)
# 4. Apply STCA modulationmodulated_spikes=self.stca(mem_potentials, spikes)
returnmodulated_spikesif__name__=="__main__":
# Setup test parametersB, T, C, H, W=4, 10, 2, 48, 48device=torch.device("cuda"iftorch.cuda.is_available() else"cpu")
# Initialize the SNN modulesnn_model=SNNInputModule(T=T, in_channels=C, out_channels=64).to(device)
test_input=torch.randn(B, T, C, H, W).to(device)
# Forward pass validationwithtorch.no_grad():
output=snn_model(test_input)
print(f"SNN STCA Module - Input Shape: {test_input.shape}")
print(f"SNN STCA Module - Output Shape: {output.shape}")

output

SNN STCA Module - Input Shape: torch.Size([4, 10, 2, 48, 48])
SNN STCA Module - Output Shape: torch.Size([4, 10, 64, 48, 48])

How to use SCAU in CNN

importtorchimporttorch.nnasnnfrommodule.SCAUimportThree_D_Coordinate_Attention_UnitclassSCAU_CNN_Block(nn.Module):
def__init__(self, channels):
super().__init__()
self.conv=nn.Conv2d(channels, channels, kernel_size=3, padding=1)
self.bn=nn.BatchNorm2d(channels)
self.relu=nn.ReLU(inplace=True)
# Initialize SCAU (3D Coordinate Attention)self.scau=Three_D_Coordinate_Attention_Unit(
inp=channels, oup=channels, reduction=16
)
defforward(self, x):
# Initial input shape: [B, C, H, W]identity=xout=self.relu(self.bn(self.conv(x)))
# Step 1: Expand to 5D [B, T, C, H, W] where T=1out=out.unsqueeze(1)
# Step 2: Apply Spatial Coordinate Attentionout=self.scau(out)
# Step 3: Squeeze back to 4D [B, C, H, W]out=out.squeeze(1)
returnout+identityif__name__=="__main__":
# Setup test parametersB, C, H, W=8, 64, 32, 32device=torch.device("cuda"iftorch.cuda.is_available() else"cpu")
# Initialize the CNN Blockcnn_block=SCAU_CNN_Block(channels=C).to(device)
test_input=torch.randn(B, C, H, W).to(device)
# Forward pass validationwithtorch.no_grad():
output=cnn_block(test_input)
print(f"CNN SCAU Block - Input Shape: {test_input.shape}")
print(f"CNN SCAU Block - Output Shape: {output.shape}")

output

CNN SCAU Block - Input Shape: torch.Size([8, 64, 32, 32])
CNN SCAU Block - Output Shape: torch.Size([8, 64, 32, 32])

If you have any questions, feel free to send an email to qingyangswu(AT)gmail DOT com. I will respond to you as soon as possible.

Thanks

We gratefully thank the authors for their wonderful works:

• Hou, Qibin, Daquan Zhou, and Jiashi Feng. "Coordinate attention for efficient mobile network design." In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 13713-13722. 2021.
• Zhu, Rui-Jie, Malu Zhang, Qihang Zhao, Haoyu Deng, Yule Duan, and Liang-Jian Deng. "TCJA-SNN: Temporal-channel joint attention for spiking neural networks." IEEE Transactions on Neural Networks and Learning Systems 36, no. 3 (2024): 5112-5125.

Reference

If you find this repo useful, please consider citing:

Yang, Qing, Meiling Zhong, Jiabin Sun, Xiurong Zhong, Shukai Duan, and Lidan Wang. "STCA-SNN: Spatio-temporal coordinate attention for spiking neural networks." Expert Systems with Applications 325 (2026): 132570. https://doi.org/10.1016/j.eswa.2026.132570

@article{yang2026stca,
title = {STCA-SNN: Spatio-Temporal Coordinate Attention for Spiking Neural Networks},
journal = {Expert Systems with Applications},
volume = {325},
pages = {132570},
year = {2026},
issn = {0957-4174},
doi = {https://doi.org/10.1016/j.eswa.2026.132570},
url = {https://www.sciencedirect.com/science/article/pii/S0957417426014831},
author={Yang, Qing and Zhong, Meiling and Sun, Jiabin and Zhong, Xiurong and Duan, Shukai and Wang, Lidan},
}